make mypy happy

src/enerbitdso/cli.py

@@ -4,7 +4,9 @@
 import operator
 import pathlib
 import sys
+from typing import Annotated, TypedDict
 
+import pydantic
 import typer
 import zoneinfo
 from rich.console import Console
@@ -17,7 +19,22 @@
 logger = logging.getLogger(__name__)
 
 DATE_FORMATS = ["%Y-%m-%d", "%Y%m%d"]
-DATE_PARTS_TO_START_DAY = {"hour": 0, "minute": 0, "second": 0, "microsecond": 0}
+
+
+class DateParts(TypedDict, total=False):
+    hour: int
+    minute: int
+    second: int
+    microsecond: int
+
+
+DATE_PARTS_TO_START_DAY: DateParts = {
+    "hour": 0,
+    "minute": 0,
+    "second": 0,
+    "microsecond": 0,
+}
+
 TZ_INFO = zoneinfo.ZoneInfo("America/Bogota")
 
 
@@ -41,9 +58,12 @@ def today():
 
 @usages.command()
 def fetch(
-    api_base_url: str = typer.Option(..., envvar="ENERBIT_API_BASE_URL"),
-    api_username: str = typer.Option(..., envvar="ENERBIT_API_USERNAME"),
-    api_password: str = typer.Option(..., envvar="ENERBIT_API_PASSWORD"),
+    api_base_url: Annotated[str, typer.Option(..., envvar="ENERBIT_API_BASE_URL")],
+    api_username: Annotated[str, typer.Option(..., envvar="ENERBIT_API_USERNAME")],
+    api_password: Annotated[
+        pydantic.SecretStr,
+        typer.Option(parser=pydantic.SecretStr, envvar="ENERBIT_API_PASSWORD"),
+    ],
     since: dt.datetime = typer.Option(
         yesterday,
         formats=DATE_FORMATS,

src/enerbitdso/enerbit.py

@@ -57,7 +57,9 @@ def get_auth_token(base_url, username, password):
     return token
 
 
-def get_client(base_url, username, password):
+def get_client(
+    base_url: str, username: str, password: pydantic.SecretStr
+) -> httpx.Client:
     url_parse: urllib.parse.ParseResult = urllib.parse.urlparse(base_url)
     url = url_parse.geturl()
     token = get_auth_token(url, username, password)
@@ -88,13 +90,16 @@ def get_schedule_usage_records(
 ) -> list[ScheduleUsageRecord]:
     path = "/measurements/schedules/usages"
     params = {
-        "since": since,
-        "until": until,
+        "since": since.isoformat(),
+        "until": until.isoformat(),
         "frt-code": frt_code,
         "period-string": "hour",
-        "period-number": 1,
+        "period-number": "1",
     }
-    response = client.get(path, params=params)
+    response = client.get(
+        path,
+        params=params,
+    )
     try:
         response.raise_for_status()
     except httpx.HTTPStatusError as e:
@@ -114,8 +119,8 @@ def get_schedule_measurement_records(
 ) -> list[ScheduleMeasurementRecord]:
     path = "/measurements/schedules/"
     params = {
-        "since": since,
-        "until": until,
+        "since": since.isoformat(),
+        "until": until.isoformat(),
         "frt-code": frt_code,
     }
     response = client.get(path, params=params)

src/enerbitdso/formats.py

@@ -1,5 +1,6 @@
 import csv
 import io
+import typing
 
 import orjson
 import pydantic
@@ -8,23 +9,22 @@
 def as_json(records: list[pydantic.BaseModel]) -> io.StringIO:
     content = orjson.dumps([r.model_dump() for r in records])
     res = io.BytesIO(content)
-    wrapper = io.TextIOWrapper(res, encoding="utf-8")
-    return wrapper
 
+    return io.StringIO(res.getvalue().decode("utf-8"))
 
-def as_csv(records: list[pydantic.BaseModel], header: bool) -> io.StringIO:
+
+def as_csv(records: typing.Sequence[pydantic.BaseModel], header: bool) -> io.StringIO:
     res = io.StringIO(newline="")
     fields = records[0].model_fields.keys()
-    content_lines = [r.model_dump() for r in records]
     writer = csv.DictWriter(res, fields, lineterminator="\n")
     if header:
         writer.writeheader()
-    for i in content_lines:
-        writer.writerow(i)
+    for i in records:
+        writer.writerow(i.model_dump())
     return res
 
 
-def as_jsonl(records: list[pydantic.BaseModel]) -> io.StringIO:
+def as_jsonl(records: typing.Sequence[pydantic.BaseModel]) -> io.StringIO:
     res = io.StringIO()
     for i in records:
         res.write(i.model_dump_json())