typing: adding types to utilities

Author: ParticularlyPythonicBS

pyproject.toml

@@ -144,7 +144,7 @@ python_version = "3.12"
 mypy_path = "stubs"
 
 # Exclude specific directories from type checking will try to add them back gradually
-exclude = "(?x)(^temoa/extensions/|^temoa/utilities/|^stubs/)"
+exclude = "(?x)(^temoa/extensions/|^stubs/)"
 
 # Strict typing for our own code
 disallow_untyped_defs = true

stubs/pyomo/core/base/component.pyi

@@ -159,6 +159,7 @@ class ComponentData(ComponentBase):
     def __idiv__(self, other: typingAny) -> typingAny: ...
     def __itruediv__(self, other: typingAny) -> typingAny: ...
     def __ipow__(self, other: typingAny) -> typingAny: ...
+    def set_value(self, val, skip_validation: bool = False) -> None: ...
 
 class ActiveComponentData(ComponentData):
     def __init__(self, component) -> None: ...

temoa/utilities/capacity_analyzer.py

@@ -5,10 +5,9 @@
 """
 
 import itertools
-import os.path
 import sqlite3
+from typing import Any
 
-from definitions import PROJECT_ROOT
 from matplotlib import pyplot as plt
 
 # Written by:  J. F. Hyink
@@ -17,34 +16,31 @@
 
 # Created on:  7/18/23
 
-# filename of db to analyze...
-db = 'US_9R_8D_CT500.sqlite'
+# filepath of db to analyze...
+source_db_file = 'US_9R_8D_CT500.sqlite'
+
 
-source_db_file = os.path.join(PROJECT_ROOT, 'data_files', 'untracked_data', db)
 print(source_db_file)
-res = []
+res: list[Any] = []
 try:
-    con = sqlite3.connect(source_db_file)
-    cur = con.cursor()
-    cur.execute('SELECT max_cap  FROM max_capacity')
-    for row in cur:
-        res.append(row)
+    with sqlite3.connect(source_db_file) as con:
+        cur: sqlite3.Cursor = con.cursor()
+        cur.execute('SELECT max_cap  FROM max_capacity')
+        for row in cur:
+            res.append(row)
 
 except sqlite3.Error as e:
     print(e)
 
-finally:
-    con.close()
-
 # chain them together into a list
-caps = list(itertools.chain(*res))
+caps: list[float] = list(itertools.chain(*res))
 
 cutoff = 1  # GW : An arbitrary cutoff between big and small capacity systems.
-small_cap_sources = [c for c in caps if c <= cutoff]
-large_cap_sources = [c for c in caps if c > cutoff]
+small_cap_sources: list[float] = [c for c in caps if c <= cutoff]
+large_cap_sources: list[float] = [c for c in caps if c > cutoff]
 
-aggregate_small_cap = sum(small_cap_sources)
-aggregate_large_cap = sum(large_cap_sources)
+aggregate_small_cap: float = sum(small_cap_sources)
+aggregate_large_cap: float = sum(large_cap_sources)
 
 print(f'{len(small_cap_sources)} small cap sources account for: {aggregate_small_cap: 0.1f} GW')
 print(f'{len(large_cap_sources)} large cap sources account for: {aggregate_large_cap: 0.1f} GW')
@@ -56,8 +52,8 @@
 # make a cumulative contribution plot, and find a 5% cutoff
 cutoff_num_sources = 0
 caps.sort()
-total_cap = sum(caps)
-cumulative_caps = [
+total_cap: float = sum(caps)
+cumulative_caps: list[float] = [
     caps[0] / total_cap,
 ]
 for i, cap in enumerate(caps[1:]):

temoa/utilities/clear_db_outputs.py

@@ -7,7 +7,7 @@
 import sys
 from pathlib import Path
 
-basic_output_tables = [
+basic_output_tables: list[str] = [
     'output_built_capacity',
     'output_cost',
     'output_curtailment',
@@ -19,15 +19,15 @@
     'output_objective',
     'output_retired_capacity',
 ]
-optional_output_tables = ['output_flow_out_summary', 'myopic_efficiency']
+optional_output_tables: list[str] = ['output_flow_out_summary', 'myopic_efficiency']
 
 if len(sys.argv) != 2:
     print('this utility file expects a CLA for the path to the database to clear')
     sys.exit(-1)
 
-target_db_str = sys.argv[1]
+target_db_str: str = sys.argv[1]
 
-proceed = input('This will clear ALL output tables in ' + target_db_str + '? (y/n): ')
+proceed: str = input('This will clear ALL output tables in ' + target_db_str + '? (y/n): ')
 if proceed == 'y':
     target_db = Path(target_db_str)
     if not target_db.exists():

temoa/utilities/database_util.py

@@ -7,10 +7,13 @@
 periods, and regions.
 """
 
+from __future__ import annotations
+
 import os
 import re
 import sqlite3
 from os import PathLike
+from typing import Any, cast
 
 import deprecated
 import pandas as pd
@@ -55,9 +58,9 @@ def __init__(self, database_path: str | PathLike[str], scenario: str | None = No
 
     def close(self) -> None:
         """Closes the database cursor and connection."""
-        if self.cur:
+        if hasattr(self, 'cur') and self.cur:
             self.cur.close()
-        if self.con:
+        if hasattr(self, 'con') and self.con:
             self.con.close()
 
     @staticmethod
@@ -115,6 +118,8 @@ def get_time_peridos_for_flags(self, flags: list[str] | None = None) -> set[int]
             query = f'SELECT period FROM time_period WHERE flag IN ({in_clause})'
 
         self.cur.execute(query)
+        # cast to int because sqlite might return strings or ints depending on how data was inserted
+        # but type hint says set[int]
         return {int(row[0]) for row in self.cur}
 
     def get_technologies_for_flags(self, flags: list[str] | None = None) -> set[str]:
@@ -125,7 +130,7 @@ def get_technologies_for_flags(self, flags: list[str] | None = None) -> set[str]
             in_clause = ', '.join(f"'{flag}'" for flag in flags)
             query = f'SELECT tech FROM Technology WHERE flag IN ({in_clause})'
 
-        return {row[0] for row in self.cur.execute(query)}
+        return {cast('str', row[0]) for row in self.cur.execute(query)}
 
     def get_commodities_and_tech(
         self, inp_comm: str | None, inp_tech: str | None, region: str | None
@@ -171,7 +176,7 @@ def get_commodities_for_flags(self, flags: list[str] | None = None) -> set[str]:
             in_clause = ', '.join(f"'{flag}'" for flag in flags)
             query = f'SELECT name FROM Commodity WHERE flag IN ({in_clause})'
 
-        return {row[0] for row in self.cur.execute(query)}
+        return {cast('str', row[0]) for row in self.cur.execute(query)}
 
     def get_commodities_by_technology(
         self, region: str | None, comm_type: str = 'input'
@@ -187,11 +192,11 @@ def get_commodities_by_technology(
         if region:
             query += f" WHERE region LIKE '%{region}%'"
 
-        return {tuple(row) for row in self.cur.execute(query)}
+        return {cast('tuple[str, str]', row) for row in self.cur.execute(query)}
 
     def get_capacity_for_tech_and_period(
         self, tech: str | None = None, period: int | None = None, region: str | None = None
-    ) -> pd.DataFrame | pd.Series:
+    ) -> pd.DataFrame | pd.Series[Any]:
         """Retrieves capacity data, aggregated by technology."""
         if not self.scenario:
             raise ValueError('A scenario must be set for output-related queries')

temoa/utilities/db_migration_to_v3.py

@@ -12,6 +12,7 @@
 import sys
 from collections import defaultdict
 from pathlib import Path
+from typing import Any
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
@@ -36,9 +37,9 @@
 new_db_name = legacy_db.stem + '_v3.sqlite'
 new_db_path = Path(legacy_db.parent, new_db_name)
 
-con_old = sqlite3.connect(legacy_db)
-con_new = sqlite3.connect(new_db_path)
-cur = con_new.cursor()
+con_old: sqlite3.Connection = sqlite3.connect(legacy_db)
+con_new: sqlite3.Connection = sqlite3.connect(new_db_path)
+cur: sqlite3.Cursor = con_new.cursor()
 
 # bring in the new schema and execute
 with open(schema_file) as src:
@@ -50,7 +51,7 @@
 
 # table mapping for DIRECT transfers
 # fmt: off
-direct_transfer_tables = [
+direct_transfer_tables: list[tuple[str, str]] = [
     ("",                    "CapacityCredit"),
     ("",                    "CapacityFactorProcess"),
     ("",                    "CapacityFactorTech"),
@@ -105,14 +106,14 @@
     ("SegFrac",             "TimeSegmentFraction"),
 ]
 
-units_added_tables = [
+units_added_tables: list[tuple[str, str]] = [
     ("",                    "MaxActivityGroup"),
     ("",                    "MaxCapacityGroup"),
     ("",                    "MinCapacityGroup"),
     ("",                    "MinActivityGroup"),
 ]
 
-sequence_added_tables = [
+sequence_added_tables: list[tuple[str, str]] = [
     ("time_season",         "TimeSeason"),
     ("time_periods",        "time_period"),
     ("time_of_day",         "TimeOfDay"),
@@ -126,12 +127,16 @@
     if old_name == '':
         old_name = new_name
 
-    new_columns = [c[1] for c in con_new.execute(f'PRAGMA table_info({new_name});').fetchall()]
-    old_columns = [c[1] for c in con_old.execute(f'PRAGMA table_info({old_name});').fetchall()]
+    new_columns: list[str] = [
+        c[1] for c in con_new.execute(f'PRAGMA table_info({new_name});').fetchall()
+    ]
+    old_columns: list[str] = [
+        c[1] for c in con_old.execute(f'PRAGMA table_info({old_name});').fetchall()
+    ]
     cols = str(old_columns[0 : len(new_columns)])[1:-1].replace("'", '')
 
     try:
-        data = con_old.execute(f'SELECT {cols} FROM {old_name}').fetchall()
+        data: list[Any] = con_old.execute(f'SELECT {cols} FROM {old_name}').fetchall()
     except sqlite3.OperationalError:
         print('TABLE NOT FOUND: ' + old_name)
         data = []
@@ -222,10 +227,11 @@
 # let's ensure all the non-global entries are consistent (same techs in each region)
 skip_rps = False
 try:
-    rps_entries = con_old.execute('SELECT * FROM tech_rps').fetchall()
+    rps_entries: list[tuple[str, str, str]] = con_old.execute('SELECT * FROM tech_rps').fetchall()
 except sqlite3.OperationalError:
     print('source does not appear to include RPS techs...skipping')
     skip_rps = True
+    rps_entries = []
 if not skip_rps:
     for region, tech, _notes in rps_entries:
         groups[region].add(tech)
@@ -239,7 +245,7 @@
     for group, techs in groups.items():
         print(f'group: {group} mismatches: {common ^ techs}')
         if group != 'global':
-            techs_common &= not common ^ techs
+            techs_common &= not (common ^ techs)
             if not techs_common:
                 print(
                     'combining RPS techs failed.  Some regions are not same.  Must be done '
@@ -357,7 +363,7 @@
 data = con_old.execute(read_qry).fetchall()
 if unlim_cap_present:
     # need to convert null -> 0 for unlim_cap to match new schema that does not allow null
-    new_data = []
+    new_data: list[Any] = []
     for row in data:
         new_row = list(row)
         if new_row[4] is None:

temoa/utilities/db_migration_v3_1_to_v4.py

@@ -16,6 +16,7 @@
 import re
 import sqlite3
 from pathlib import Path
+from typing import Any
 
 # ---------- Mapping configuration ----------
 CUSTOM_MAP: dict[str, str] = {
@@ -93,7 +94,7 @@ def map_token_no_cascade(token: str) -> str:
     return to_snake_case(token)
 
 
-def get_table_info(conn: sqlite3.Connection, table: str) -> list[tuple]:
+def get_table_info(conn: sqlite3.Connection, table: str) -> list[tuple[Any, ...]]:
     try:
         return conn.execute(f'PRAGMA table_info({table});').fetchall()
     except sqlite3.OperationalError:
@@ -131,7 +132,7 @@ def migrate_direct_table(
     return len(filtered)
 
 
-def migrate_all(args) -> None:
+def migrate_all(args: argparse.Namespace) -> None:
     src = Path(args.source)
     schema = Path(args.schema)
     out = Path(args.out) if args.out else src.with_suffix('.v4.sqlite')