Ruff fixes (various message classes per commit)

create_trace_mapping.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import yaml
 from generator_to_trace_draft_mapper import (
     draft_solar_generator_to_trace_mapping,
@@ -19,14 +21,14 @@
 solar_generator_mapping = draft_solar_generator_to_trace_mapping(
     solar_gens, solar_traces
 )
-with open("draft_solar_generator_mapping.yaml", "w") as file:
+with Path("draft_solar_generator_mapping.yaml").open("w") as file:
     yaml.dump(solar_generator_mapping, file, default_flow_style=False)
 
 
 solar_traces = "/media/nick/Samsung_T5/isp_2024_data/trace_data/solar/solar_2023"
 rezs = gets_rezs(workbook)
 solar_rez_mapping = draft_solar_rez_mapping(rezs, solar_traces)
-with open("solar_area_mapping.yaml", "w") as file:
+with Path("solar_area_mapping.yaml").open("w") as file:
     yaml.dump(solar_rez_mapping, file, default_flow_style=False)
 
 duids_and_station_names = static_table(
@@ -47,12 +49,12 @@
 wind_generator_mapping = draft_wind_generator_to_trace_mapping(
     wind_gens, wind_duids_and_station_names, wind_traces
 )
-with open("draft_wind_generator_mapping.yaml", "w") as file:
+with Path("draft_wind_generator_mapping.yaml").open("w") as file:
     yaml.dump(wind_generator_mapping, file, default_flow_style=False, sort_keys=False)
 
 
 wind_traces = "D:/isp_2024_data/trace_data/wind/wind_2023"
 rezs = gets_rezs(workbook)
 wind_rez_mapping = draft_wind_rez_mapping(rezs, wind_traces)
-with open("draft_wind_rez_mapping.yaml", "w") as file:
+with Path("draft_wind_rez_mapping.yaml").open("w") as file:
     yaml.dump(wind_rez_mapping, file, default_flow_style=False)

src/isp_trace_parser/__init__.py

@@ -15,13 +15,13 @@
 from isp_trace_parser.wind_traces import WindMetadataFilter, parse_wind_traces
 
 __all__ = [
-    "trace_formatter",
+    "DemandMetadataFilter",
+    "SolarMetadataFilter",
+    "WindMetadataFilter",
+    "construct_reference_year_mapping",
     "get_data",
-    "parse_wind_traces",
     "parse_demand_traces",
     "parse_solar_traces",
-    "construct_reference_year_mapping",
-    "WindMetadataFilter",
-    "SolarMetadataFilter",
-    "DemandMetadataFilter",
+    "parse_wind_traces",
+    "trace_formatter",
 ]

src/isp_trace_parser/construct_reference_year_mapping.py

@@ -13,7 +13,7 @@
 @validate_call
 def construct_reference_year_mapping(
     start_year: int, end_year: int, reference_years: list[int]
-):
+) -> dict:
     """Constructs a dictionary mapping a sequence of modeling years to a cycle of reference years.
 
     Examples:
@@ -42,4 +42,4 @@ def construct_reference_year_mapping(
         reference_years = (
             reference_years * full_reference_year_cycles
         ) + reference_years[:partial_cycle_length]
-    return dict(zip(years, reference_years))
+    return dict(zip(years, reference_years, strict=True))

src/isp_trace_parser/demand_trace_metadata.py

@@ -29,7 +29,8 @@ def build(
         refyear, _, dimensions_suffix = after.partition("_")
         key = (location_prefix, dimensions_suffix)
         if not refyear.isdigit() or key not in lookup:
-            raise ValueError(f"Unexpected trace filename: {path.name}")
+            msg = f"Unexpected trace filename: {path.name}"
+            raise ValueError(msg)
         file_metadata[path] = {**lookup[key], "reference_year": int(refyear)}
     return file_metadata
 

src/isp_trace_parser/demand_traces.py

@@ -8,7 +8,7 @@
 import functools
 import os
 from pathlib import Path
-from typing import Literal, Optional
+from typing import Literal
 
 import polars as pl
 from joblib import Parallel, delayed
@@ -52,15 +52,16 @@ class DemandMetadataFilter(BaseModel):
         reference_year: list of ints specifying reference_years
     """
 
-    subregion: Optional[list[str]] = None
-    scenario: Optional[
+    subregion: list[str] | None = None
+    scenario: (
         list[Literal["Step Change", "Progressive Change", "Green Energy Exports"]]
-    ] = None
-    poe: Optional[list[Literal["POE50", "POE10"]]] = None
-    demand_type: Optional[
-        list[Literal["OPSO_MODELLING", "OPSO_MODELLING_PVLITE", "PV_TOT"]]
-    ] = None
-    reference_year: Optional[list[int]] = None
+        | None
+    ) = None
+    poe: list[Literal["POE50", "POE10"]] | None = None
+    demand_type: (
+        list[Literal["OPSO_MODELLING", "OPSO_MODELLING_PVLITE", "PV_TOT"]] | None
+    ) = None
+    reference_year: list[int] | None = None
 
 
 @validate_call
@@ -69,7 +70,7 @@ def parse_demand_traces(
     parsed_directory: str | Path,
     use_concurrency: bool = True,
     filters: DemandMetadataFilter | None = None,
-):
+) -> None:
     """Takes a directory with AEMO demand trace data and reformats the data, saving it to a new directory.
 
     AEMO demand trace data comes in CSVs with columns specifying the year, day, and month, and data columns

src/isp_trace_parser/get_data.py

@@ -7,7 +7,7 @@
 
 import datetime
 from pathlib import Path
-from typing import List, Literal
+from typing import Literal
 
 import pandas as pd
 import polars as pl
@@ -16,7 +16,7 @@
 
 def _year_range_to_dt_range(
     start_year: int, end_year: int, year_type: Literal["fy", "calendar"] = "fy"
-):
+) -> datetime.datetime:
     """
     Convert year range to datetime boundaries for efficient time filtering.
 
@@ -44,19 +44,20 @@ def _year_range_to_dt_range(
             end_year, 7, 1
         )
 
-    elif year_type == "calendar":
+    if year_type == "calendar":
         return datetime.datetime(start_year, 1, 1), datetime.datetime(
             end_year + 1, 1, 1
         )
+    raise ValueError(year_type)
 
 
 def _query_parquet_single_reference_year(
     start_year: int,
     end_year: int,
     reference_year: int,
     directory: str | Path,
-    filters: dict[str, any] = None,
-    select_columns: list[str] = None,
+    filters: dict[str, any] | None = None,
+    select_columns: list[str] | None = None,
     year_type: Literal["fy", "calendar"] = "fy",
 ) -> pd.DataFrame:
     """
@@ -143,20 +144,19 @@ def _query_parquet_multiple_reference_years(
                 start_year=year, end_year=year, reference_year=reference_year, **kwargs
             )
         )
-    data = pd.concat(data).reset_index(drop=True)
-    return data
+    return pd.concat(data).reset_index(drop=True)
 
 
 @validate_call
 def get_project_single_reference_year(
     start_year: int,
     end_year: int,
     reference_year: int,
-    project: str | List,
+    project: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query project trace data for a single reference year.
 
@@ -244,12 +244,12 @@ def get_zone_single_reference_year(
     start_year: int,
     end_year: int,
     reference_year: int,
-    zone: str | List,
-    resource_type: str | List,
+    zone: str | list,
+    resource_type: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query zone trace data for a single reference year.
 
@@ -340,14 +340,14 @@ def get_demand_single_reference_year(
     start_year: int,
     end_year: int,
     reference_year: int,
-    scenario: str | List,
-    subregion: str | List,
-    demand_type: str | List,
-    poe: str | List,
+    scenario: str | list,
+    subregion: str | list,
+    demand_type: str | list,
+    poe: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query demand trace data for a single reference year.
 
@@ -448,11 +448,11 @@ def get_demand_single_reference_year(
 @validate_call
 def get_project_multiple_reference_years(
     reference_year_mapping: dict[int, int],
-    project: str | List,
+    project: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query project trace data across multiple reference years.
 
@@ -537,12 +537,12 @@ def get_project_multiple_reference_years(
 @validate_call
 def get_zone_multiple_reference_years(
     reference_year_mapping: dict[int, int],
-    zone: str | List,
-    resource_type: str | List,
+    zone: str | list,
+    resource_type: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query zone trace data across multiple reference years.
 
@@ -630,14 +630,14 @@ def get_zone_multiple_reference_years(
 @validate_call
 def get_demand_multiple_reference_years(
     reference_year_mapping: dict[int, int],
-    scenario: str | List,
-    subregion: str | List,
-    demand_type: str | List,
-    poe: str | List,
+    scenario: str | list,
+    subregion: str | list,
+    demand_type: str | list,
+    poe: str | list,
     directory: str | Path,
     year_type: Literal["fy", "calendar"] = "fy",
-    select_columns: list[str] = None,
-):
+    select_columns: list[str] | None = None,
+) -> pd.DataFrame:
     """
     Query demand trace data across multiple reference years.
 

src/isp_trace_parser/input_validation.py

@@ -11,7 +11,7 @@
 def input_directory(path: Path | str) -> Path:
     path = is_valid_path(path)
     if not path.is_dir():
-        raise ValueError(f"Directory {path} does not exist")
+        raise FileNotFoundError(path)
     return path
 
 
@@ -22,10 +22,12 @@ def parsed_directory(path: str | Path) -> Path:
 def is_valid_path(path: str | Path) -> Path:
     try:
         return Path(path)
-    except (TypeError, ValueError):
-        raise ValueError(f"Invalid parsed directory path: {path}")
+    except (TypeError, ValueError) as exc:
+        msg = f"Invalid parsed directory path: {path}"
+        raise ValueError(msg) from exc
 
 
-def start_year_before_end_year(start_year, end_year):
+def start_year_before_end_year(start_year: int, end_year: int) -> None:
     if end_year < start_year:
-        raise ValueError(f"Start year {end_year} < end year {start_year}")
+        msg = f"Start year {end_year} < end year {start_year}"
+        raise ValueError(msg)

src/isp_trace_parser/optimise_parquet.py

@@ -7,7 +7,6 @@
 
 from itertools import product
 from pathlib import Path
-from typing import Optional
 
 import duckdb
 from pydantic import validate_call
@@ -30,7 +29,7 @@ def partition_traces_by_columns(
     input_directory: str | Path,
     output_directory: str | Path,
     partition_cols: list[str],
-    sort_by: Optional[list[str]] = ["datetime"],
+    sort_by: list[str] | None = None,
 ) -> None:
     """Partition parquet traces by specified columns with optional sorting.
 
@@ -60,6 +59,12 @@ def partition_traces_by_columns(
         ...     partition_cols=["scenario", "reference_year"]
         ... ) # doctest: +SKIP
     """
+
+    if sort_by is None:
+        # Avoid use of mutable data structure for argument defaults
+        # (see Ruff rule B006).
+        sort_by = ["datetime"]
+
     output_path = Path(output_directory)
     output_path.mkdir(parents=True, exist_ok=True)
 
@@ -71,23 +76,21 @@ def partition_traces_by_columns(
         values = con.execute(f"""
             SELECT DISTINCT {col}
             FROM read_parquet('{input_directory}')
-        """).fetchall()
+        """).fetchall()  # noqa: S608
         distinct_values.append(values)
 
     partitions = [tuple(val[0] for val in vals) for vals in product(*distinct_values)]
 
     for partition_values in partitions:
-        # print(*partition_values)
-
         conditions = []
-        for col, val in zip(partition_cols, partition_values):
+        for col, val in zip(partition_cols, partition_values, strict=True):
             if isinstance(val, str):
                 conditions.append(f"{col}='{val}'")
             else:
                 conditions.append(f"{col}={val}")
 
         where_clause = " AND ".join(conditions)
-        query = f"SELECT * FROM read_parquet('{input_directory}') WHERE {where_clause}"
+        query = f"SELECT * FROM read_parquet('{input_directory}') WHERE {where_clause}"  # noqa: S608
 
         if sort_by:
             query += f" ORDER BY {', '.join(sort_by)}"