Typing fixes

.github/workflows/python-package-main.yml

@@ -24,7 +24,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.12", "3.13"]
       fail-fast: false
     defaults:
       run:
@@ -42,7 +42,7 @@ jobs:
         run: uv run python -m pytest tests
 
       - name: Check typing
-        run: uv run ty check
+        run: uv run basedpyright
 
   check-docs:
     runs-on: ubuntu-latest

pyproject.toml

@@ -44,13 +44,13 @@ tariff-fetch-gas = "tariff_fetch.cli_gas:main_cli"
 dev = [
   "pytest>=7.2.0",
   "ruff>=0.11.5",
-  "ty>=0.0.1a21",
   "deptry>=0.20.0",
   "mkdocs>=1.6.0",
   "mkdocs-material>=9.5.0",
   "mkdocstrings[python]>=0.26.1",
   "build>=1.0.0",
   "twine>=4.0.0",
+  "basedpyright>=1.38.1",
 ]
 [build-system]
 requires = ["hatchling"]
@@ -75,6 +75,6 @@ ignore = ["E501"]
 [tool.ruff.lint.per-file-ignores]
 "tests/*" = ["S101"]
 
-[tool.ty.environment]
-python = "./.venv"
-python-version = "3.13"
+[tool.basedpyright]
+include = ["tariff_fetch"]
+pythonVersion = "3.11"

tariff_fetch/_cli/__init__.py

@@ -1,6 +1,7 @@
 import os
 from datetime import datetime
 from pathlib import Path
+from typing import cast
 
 import questionary
 from pathvalidate import sanitize_filename
@@ -24,10 +25,13 @@ def prompt_filename(output_folder: Path, suggested_filename: str, extension: str
         filepath = output_folder.joinpath(f"{suggested_filename}-0{os.extsep}{extension}")
 
     return Path(
-        questionary.path(
-            message="Path to save the results",
-            default=filepath.as_posix(),
-            file_filter=lambda _: Path(_).suffix == extension,
-            validate=lambda _: (not os.path.exists(_)) or "A file with that name already exists",
-        ).ask()
+        cast(
+            str,
+            questionary.path(
+                message="Path to save the results",
+                default=filepath.as_posix(),
+                file_filter=lambda _: Path(_).suffix == extension,
+                validate=lambda _: (not os.path.exists(_)) or "A file with that name already exists",  # pyright: ignore[reportUnknownLambdaType, reportUnknownArgumentType]
+            ).ask(),
+        )
     )

tariff_fetch/_cli/genability.py

@@ -1,26 +1,30 @@
-import json
 import os
-from datetime import datetime, timezone
+from datetime import date
 from pathlib import Path
+from typing import cast
 
 import questionary
 from dotenv import load_dotenv
+from pydantic import TypeAdapter
 
-from tariff_fetch.genability.lse import get_lses_page
-from tariff_fetch.genability.tariffs import CustomerClass, TariffType, tariffs_paginate
+# from tariff_fetch.genability.lse import get_lses_page
+# from tariff_fetch.genability.tariffs import CustomerClass, TariffType, tariffs_paginate
+from tariff_fetch.arcadia.api import ArcadiaSignalAPI
+from tariff_fetch.arcadia.schema import tariff
+from tariff_fetch.arcadia.schema.common import CustomerClass, TariffType
 
 from . import console, prompt_filename
 from .types import Utility
 
 
-def _find_utility_lse_id(utility: Utility) -> int | None:
+def _find_utility_lse_id(api: ArcadiaSignalAPI, utility: Utility) -> int | None:
     with console.status("Fetching lses..."):
-        lses = get_lses_page(
+        lses = api.lses.get_page(
             fields="min",
-            searchOn=["code"],
+            search_on=["code"],
             search=str(utility.eia_id),
-            startsWith=True,
-            endsWith=True,
+            starts_with=True,
+            ends_with=True,
         )["results"]
     if len(lses) == 0:
         # No utilities found with this eia id
@@ -30,103 +34,117 @@ def _find_utility_lse_id(utility: Utility) -> int | None:
         return None
     if len(lses) == 1:
         # Found one utility
-        utility_lse_id = lses[0]["lseId"]
+        utility_lse_id = lses[0]["lse_id"]
         return utility_lse_id
     else:
         # Nothing found; this should *theoretically* never happen but let's keep it just in case
-        choices = [questionary.Choice(title=_["name"], value=_["lseId"]) for _ in lses]
+        choices = [questionary.Choice(title=_["name"], value=_["lse_id"]) for _ in lses]
         choices.append(questionary.Separator())
         choices.append(questionary.Choice(title="None of these", value=None))
-        utility_lse_id: int | None = questionary.select(
-            message=f"Found multiple utilities with lse id = {utility.eia_id}. Select one.", choices=choices
-        ).ask()
+        utility_lse_id = cast(
+            int | None,
+            questionary.select(
+                message=f"Found multiple utilities with lse id = {utility.eia_id}. Select one.", choices=choices
+            ).ask(),
+        )
         if utility_lse_id is None:
             console.print("No utility chosen")
             return None
         return utility_lse_id
 
 
 def _select_tariffs(
-    lse_id: int, customer_classes: list[CustomerClass], tariff_types: list[TariffType]
+    api: ArcadiaSignalAPI, lse_id: int, customer_classes: list[CustomerClass], tariff_types: list[TariffType]
 ) -> list[tuple[str, int]]:
     with console.status("Fetching tariffs..."):
         tariffs = list(
-            tariffs_paginate(
-                lseId=lse_id,
-                fields="min",
-                effectiveOn=datetime.now(timezone.utc),
-                customerClasses=customer_classes,
-                tariffTypes=tariff_types,
+            api.tariffs.iter_pages(
+                lse_id=lse_id,
+                effective_on=date.today(),
+                customer_classes=customer_classes,
+                tariff_types=tariff_types,
             )
         )
     if not tariffs:
         return []
-    return questionary.checkbox(
-        message="Select tariffs",
-        choices=[
-            questionary.Choice(title=_["tariffName"], value=(_["tariffName"], _["masterTariffId"]), checked=True)
-            for _ in tariffs
-        ],
-        use_search_filter=True,
-        use_jk_keys=False,
-    ).ask()
+    return cast(
+        list[tuple[str, int]],
+        questionary.checkbox(
+            message="Select tariffs",
+            choices=[
+                questionary.Choice(
+                    title=f"{_['tariff_name']} ({_['tariff_id']})",
+                    value=(_["tariff_name"], _["master_tariff_id"]),  # pyright: ignore[reportAny]
+                    checked=True,
+                )
+                for _ in tariffs
+            ],
+            use_search_filter=True,
+            use_jk_keys=False,
+        ).ask(),
+    )
 
 
 def _select_customer_classes() -> list[CustomerClass]:
-    return questionary.checkbox(
-        message="Select customer classes",
-        choices=[
-            questionary.Choice(title="Residential", value="RESIDENTIAL"),
-            questionary.Choice(title="General", value="GENERAL"),
-            questionary.Choice(title="Special Use", value="SPECIAL_USE"),
-        ],
-        validate=lambda _: True if _ else "Select at least one customer class",
-    ).ask()
+    return cast(
+        list[CustomerClass],
+        questionary.checkbox(
+            message="Select customer classes",
+            choices=[
+                questionary.Choice(title="Residential", value="RESIDENTIAL"),
+                questionary.Choice(title="General", value="GENERAL"),
+                questionary.Choice(title="Special Use", value="SPECIAL_USE"),
+            ],
+            validate=lambda _: True if _ else "Select at least one customer class",
+        ).ask(),
+    )
 
 
 def _select_tariff_types() -> list[TariffType]:
-    return questionary.checkbox(
-        message="Select tariff types",
-        choices=[
-            questionary.Choice(title="Default", value="DEFAULT"),
-            questionary.Choice(title="Alternative", value="ALTERNATIVE"),
-            questionary.Choice(title="Optional extra", value="OPTIONAL_EXTRA"),
-            questionary.Choice(title="Rider", value="RIDER"),
-        ],
-        validate=lambda _: bool(_) or "Select at least one tariff type",
-    ).ask()
-
-
-def _fetch_tariffs(tariffs: list[tuple[str, int]]):
-    result = []
+    return cast(
+        list[TariffType],
+        questionary.checkbox(
+            message="Select tariff types",
+            choices=[
+                questionary.Choice(title="Default", value="DEFAULT"),
+                questionary.Choice(title="Alternative", value="ALTERNATIVE"),
+                questionary.Choice(title="Optional extra", value="OPTIONAL_EXTRA"),
+                questionary.Choice(title="Rider", value="RIDER"),
+            ],
+            validate=lambda _: bool(_) or "Select at least one tariff type",
+        ).ask(),
+    )
+
+
+def _fetch_tariffs(api: ArcadiaSignalAPI, tariffs: list[tuple[str, int]]):
+    result: list[tariff.TariffExtended] = []
     with console.status("Fetching tariffs..."):
         for name, id_ in tariffs:
             console.print(f"Tariff id: {name}")
-            page = list(
-                tariffs_paginate(
-                    masterTariffId=id_,
-                    effectiveOn=datetime.now(timezone.utc),
-                    fields="ext",
-                    populateProperties=True,
-                    populateRates=True,
-                )
+            page = api.tariffs.iter_pages(
+                fields="ext",
+                master_tariff_id=id_,
+                effective_on=date.today(),
+                populate_properties=True,
+                populate_rates=True,
             )
             result.extend(page)
     return result
 
 
 def process_genability(utility: Utility, output_folder: Path):
-    load_dotenv()
+    _ = load_dotenv()
     if not os.getenv("ARCADIA_APP_ID"):
         console.print("[b]ARCADIA_APP_ID[/] environment variable is not set.")
     if not os.getenv("ARCADIA_APP_KEY"):
         console.print("[b]ARCADIA_APP_KEY[/] environment variable is not set.")
     if not (os.getenv("ARCADIA_APP_ID") and os.getenv("ARCADIA_APP_KEY")):
         console.print("Cannot use Arcadia API due to missing credentials")
-        console.input("Press enter to proceed...")
+        _ = console.input("Press enter to proceed...")
         return
+    api = ArcadiaSignalAPI()
 
-    lse_id = _find_utility_lse_id(utility)
+    lse_id = _find_utility_lse_id(api, utility)
     if lse_id is None:
         return
 
@@ -136,17 +154,17 @@ def process_genability(utility: Utility, output_folder: Path):
     if not (tariff_types := _select_tariff_types()):
         return
 
-    if not (tariffs := _select_tariffs(lse_id, customer_classes, tariff_types)):
+    if not (tariffs := _select_tariffs(api, lse_id, customer_classes, tariff_types)):
         console.print("[red]No tariffs found[/]")
-        console.input("Press enter to proceed...")
+        _ = console.input("Press enter to proceed...")
         return
 
-    results = _fetch_tariffs(tariffs)
+    results = _fetch_tariffs(api, tariffs)
     suggested_filename = f"arcadia_{utility.name}"
 
     if not (filename := prompt_filename(output_folder, suggested_filename, "json")):
         return
 
     filename.parent.mkdir(exist_ok=True)
-    filename.write_text(json.dumps(results, indent=2))
+    _ = filename.write_bytes(TypeAdapter(list[tariff.TariffExtended]).dump_json(results, indent=2))
     console.print(f"Wrote [blue]{len(results)}[/] records to {filename}")