chore: update ty to 0.0.44

.mise.toml

@@ -24,7 +24,7 @@ ripgrep = "latest"
 dprint = "latest"
 uv = "0.9.30"
 ruff = "0.15.0"
-ty = "0.0.32"
+ty = "0.0.44"
 
 "aqua:j178/prek" = "latest"
 

STYLE_GUIDE.md

@@ -141,6 +141,18 @@ def process(data: pd.DataFrame, columns: Sequence[str] | None = None) -> Self:
 - `Protocol` for structural subtyping when you need duck-typing boundaries
 - Avoid `Any` -- prefer `object`, generics, or `Protocol`
 - `TYPE_CHECKING` guards for heavy imports (`pandas`, `torch`, `transformers`); not needed for stdlib or lightweight imports
+- `TypeIs[T]` from `typing_extensions` for side-effect-free predicates that fully validate a value as `T`. Use `bool` for ordinary predicates that do not establish a type, and reserve `TypeGuard` for cases `TypeIs` cannot express because the narrowed type is not compatible with the input type.
+
+```python
+from typing_extensions import TypeIs
+
+
+def is_json_object(value: object) -> TypeIs[dict[str, JsonValue]]:
+    return isinstance(value, dict) and all(
+        isinstance(key, str) and is_json_value(item)
+        for key, item in value.items()
+    )
+```
 
 Legacy modules (`pii_replacer/`) still use `Optional`/`List`/`Dict`. Only `pii_replacer/` is excluded from `ty` type-checking (see `[tool.ty.src] exclude` in `pyproject.toml`). The remaining legacy usages will be migrated when those modules come under the type checker.
 

src/nemo_safe_synthesizer/cli/datasets.py

@@ -21,6 +21,22 @@
 logger = get_logger(__name__)
 
 
+def _require_dataframe(value: object, *, url: str) -> pd.DataFrame:
+    if isinstance(value, pd.DataFrame):
+        return value
+    raise TypeError(f"Expected dataset reader for {url} to return a pandas DataFrame, got {type(value).__name__}")
+
+
+def _dynamic_callable(value: object) -> Any:
+    """Widen a typed pandas reader callable to ``Any``.
+
+    ty raises ``invalid-argument-type`` when a union of overloaded pandas reader
+    callables is dispatched with ``**kwargs``. Erasing the callable type here
+    keeps the runtime behavior unchanged while avoiding that diagnostic.
+    """
+    return value
+
+
 class DatasetInfo(BaseModel):
     """Entry in the dataset registry."""
 
@@ -96,20 +112,22 @@ def fetch(self) -> pd.DataFrame:
         logger.info(f"Reading dataset from {url}")
 
         # Determine the file extension and appropriate reader
-        match Path(url).suffix.lstrip("."):
+        reader: Any
+        extension = Path(url).suffix.lstrip(".")
+        match extension:
             case "csv" | "txt":
-                reader = pd.read_csv
+                reader = _dynamic_callable(pd.read_csv)
                 default_load_args: dict[str, Any] = {}
             case "json":
-                reader = pd.read_json
+                reader = _dynamic_callable(pd.read_json)
                 default_load_args = {}
             case "jsonl":
-                reader = pd.read_json
+                reader = _dynamic_callable(pd.read_json)
                 default_load_args = {"lines": True}
             case "parquet":
-                reader = pd.read_parquet
+                reader = _dynamic_callable(pd.read_parquet)
                 default_load_args = {}
-            case extension:
+            case _:
                 if not extension:
                     extension = f"<no extension found on url '{url}'>"
                 raise ValueError(f"Unsupported file extension: {extension}")
@@ -118,7 +136,7 @@ def fetch(self) -> pd.DataFrame:
         final_load_args = {**default_load_args, **(self.load_args or {})}
 
         try:
-            return reader(url, **final_load_args)  # ty: ignore[invalid-argument-type] -- reader union includes parquet which has stricter signature
+            return _require_dataframe(reader(url, **final_load_args), url=url)
         except Exception as e:
             logger.error(f"Error reading dataset from {url}: {e}", exc_info=True)
             raise

src/nemo_safe_synthesizer/data_processing/actions/dates.py

@@ -13,7 +13,7 @@
 import itertools
 import re
 from collections import Counter
-from collections.abc import Iterable, Iterator
+from collections.abc import Hashable, Iterable, Iterator
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from random import randint
@@ -399,7 +399,7 @@ def infer_from_series(date_series: Iterable[str]) -> Optional[str]:
 def fit_and_transform_dates(
     df: pd.DataFrame,
     inplace: bool = False,
-) -> tuple[dict[str, dict[str, str]], pd.DataFrame]:
+) -> tuple[dict[Hashable, dict[str, str]], pd.DataFrame]:
     """Detect date columns, convert them to elapsed seconds, and record the transformation.
 
     For each object-typed column, samples values to infer a date format. If
@@ -412,10 +412,10 @@ def fit_and_transform_dates(
 
     Returns:
         A tuple of (date_min_dict, result_df). ``date_min_dict`` maps column
-        names to ``{"format": ..., "min": ...}`` dicts needed by
+        labels to ``{"format": ..., "min": ...}`` dicts needed by
         ``transform_dates`` for reversal.
     """
-    date_min_dict = {}
+    date_min_dict: dict[Hashable, dict[str, str]] = {}
     object_cols = [col for col, col_type in df.dtypes.items() if col_type == "object"]
     result_df = df.copy() if not inplace else df
     for object_col in object_cols:
@@ -437,11 +437,11 @@ def fit_and_transform_dates(
     return date_min_dict, result_df
 
 
-def transform_dates(dates: dict[str, dict[str, str]], df: pd.DataFrame) -> pd.DataFrame:
+def transform_dates(dates: dict[Hashable, dict[str, str]], df: pd.DataFrame) -> pd.DataFrame:
     """Apply a previously fitted date-to-seconds transformation to a DataFrame.
 
     Args:
-        dates: Mapping from column names to ``{"format": ..., "min": ...}``
+        dates: Mapping from column labels to ``{"format": ..., "min": ...}``
             dicts as returned by ``fit_and_transform_dates``.
         df: DataFrame to transform.
 

src/nemo_safe_synthesizer/data_processing/actions/distributions.py

@@ -13,7 +13,6 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from functools import partial
 from typing import Annotated, Any, Literal, Optional, Union
 
 import numpy as np
@@ -77,21 +76,20 @@ def _round_datetime(self, dt: datetime, precision: timedelta) -> datetime:
         rounded_ts = round(dt.timestamp() / precision.total_seconds()) * precision.total_seconds()
         return datetime.fromtimestamp(rounded_ts)
 
-    def _apply_universal_params(self, samples: list[datetime]) -> list[datetime]:
-        ret: list[str] | list[datetime] = samples
-
-        ops = []
-        if self.precision is not None:
-            ops.append(partial(self._round_datetime, self.precision))
+    def _apply_universal_params(self, samples: list[datetime]) -> list[datetime] | list[str]:
         if self.format is not None:
-            ops.append(lambda x: x.strftime(self.format))
-
+            formatted: list[str] = []
+            for sample in samples:
+                if self.precision is not None:
+                    sample = self._round_datetime(sample, self.precision)
+                formatted.append(sample.strftime(self.format))
+            return formatted
+
+        ret: list[datetime] = []
         for sample in samples:
-            n = sample
-            for op in ops:
-                n = op(n)
-            ret.append(n)
-
+            if self.precision is not None:
+                sample = self._round_datetime(sample, self.precision)
+            ret.append(sample)
         return ret
 
 

src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py

@@ -6,6 +6,7 @@
 import itertools
 import math
 import re
+from collections.abc import Hashable
 from datetime import datetime
 from decimal import Decimal
 from functools import cached_property
@@ -138,14 +139,15 @@ def _pandas_entropy(column: pd.Series, base: float | None = None) -> np.float64:
         return -(vc * np.log(vc) / np.log(base)).sum()
 
     @staticmethod
-    def _is_really_categorical(column: str) -> bool:
+    def _is_really_categorical(column: Hashable) -> bool:
+        column_name = str(column)
         # Break the header up into parts
         for separator in ["_", " ", "-", "."]:
-            if separator in column:
-                col_parts = column.split(separator)
+            if separator in column_name:
+                col_parts = column_name.split(separator)
                 break
             else:
-                col_parts = [column]
+                col_parts = [column_name]
 
         # Go through the parts and divide up camel case
         col_final_parts = []
@@ -366,16 +368,16 @@ def _aia(
 
             # Get all combinations of columns to be the quasi-identifiers
             # This gets explosive when column count > 500
+            training_columns = list(training_df.columns)
             if len(training_df.columns) < 500:
-                qi_combos = list(itertools.combinations(training_df.columns, quasi_identifier_count))
+                qi_combos = list(itertools.combinations(training_columns, quasi_identifier_count))
             else:
-                columns = list(training_df.columns)
                 qi_combos = []
-                for i in range(len(columns) - quasi_identifier_count):
-                    combo = set()
+                for i in range(len(training_columns) - quasi_identifier_count + 1):
+                    combo = []
                     for j in range(quasi_identifier_count):
-                        combo.add(columns[i + j])
-                    qi_combos.append(combo)
+                        combo.append(training_columns[i + j])
+                    qi_combos.append(tuple(combo))
 
             np.random.seed(5)
             np.random.shuffle(qi_combos)
@@ -415,7 +417,6 @@ def _aia(
 
             # As we process the attack dataset, we'll accumulate for each column the number of
             # correct and incorrect predictions
-            training_columns = [str(column) for column in training_df.columns]
             correct = {predict_column: 0 for predict_column in training_columns}
             incorrect = {predict_column: 0 for predict_column in training_columns}
 
@@ -431,14 +432,15 @@ def _aia(
                     more_to_process = False
                     continue
 
-                # Randomly sample columns to be the `qi`
-                qi = qi_combos[qi_index]
-                qi_index += 1
                 # We stop processing if all qi combos have been processed.
                 if qi_index == len(qi_combos):
                     more_to_process = False
                     continue
 
+                # Randomly sample columns to be the `qi`
+                qi = qi_combos[qi_index]
+                qi_index += 1
+
                 # Predict columns are all but the `qi`
                 predict_columns = [column for column in training_columns if column not in qi]
 
@@ -578,7 +580,8 @@ def _aia(
                 for i in range(len(entropy)):
                     entropy_wts.append(0)
             else:
-                arr = (entropy - min(entropy)) / (max(entropy) - min(entropy))
+                entropy_arr = np.asarray(entropy, dtype=float)
+                arr = (entropy_arr - min(entropy)) / (max(entropy) - min(entropy))
                 entropy_wts = arr / arr.sum()
 
             i = 0

src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py

@@ -248,8 +248,8 @@ def _compute_mia(
     ) -> tuple[
         float,
         list[str],
-        dict[str, list[int]],
-        dict[str, list[int]],
+        dict[float, int],
+        dict[float, int],
     ]:
         """Core membership inference attack implementation for a single run.
 
@@ -286,8 +286,8 @@ def _compute_mia(
             pd.concat([training_df_attack, test_df_norm]).reset_index(drop=True).sample(frac=1, random_state=run)
         )
 
-        attack_synth_dist_text = [[0] for i in range(len(attack_df))]
-        attack_synth_indices_text = [[0] for i in range(len(attack_df))]
+        attack_synth_dist_text: list[list[float]] = [[0.0] for i in range(len(attack_df))]
+        attack_synth_indices_text: list[list[int]] = [[0] for i in range(len(attack_df))]
 
         # Get the NN dist for text for the entire attack dataset
 
@@ -362,8 +362,8 @@ def _compute_mia(
         score = 0
         attack_summary = []
 
-        tp_cnts = {}
-        fp_cnts = {}
+        tp_cnts: dict[float, int] = {}
+        fp_cnts: dict[float, int] = {}
 
         # Using the above text and tabular distances we now compute an overall distance score for
         # every record in the attack dataset. We then conduct 36 individual mia attacks on this one big
@@ -519,8 +519,8 @@ def mia(
 
             scores = []
             attack_sum_values = []
-            tps_values = {}
-            fps_values = {}
+            tps_values: dict[float, int] = {}
+            fps_values: dict[float, int] = {}
             for i in [0.1, 0.2, 0.3, 0.4]:
                 tps_values[i] = 0
                 fps_values[i] = 0

src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py

@@ -676,15 +676,16 @@ def generate_text_structure_similarity_figures(
         "average_words_per_sentence",
         "average_characters_per_word",
     ]
-    figures = []
+    figures: list[go.Figure] = []
     for key in statistics_keys:
         if training_statistics.per_record_statistics.empty or synthetic_statistics.per_record_statistics.empty:
             break
         figure = histogram_figure(
             training_statistics.per_record_statistics[key],
             synthetic_statistics.per_record_statistics[key],
         )
-        figures.append(figure)
+        if figure is not None:
+            figures.append(figure)
     if not figures:
         return None
 

src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py

@@ -21,6 +21,7 @@
     ColumnDistribution,
     ColumnDistributionPlotRow,
 )
+from ....evaluation.components.component import Component
 from ....evaluation.components.correlation import (
     Correlation,
 )
@@ -150,7 +151,7 @@ def from_dataframes(
             mandatory_columns=MultimodalReport._get_config_value("mandatory_columns", [], config),
         )
 
-        components = []
+        components: list[Component] = []
 
         attribute_inference_protection = AttributeInferenceProtection(
             score=EvaluationScore(grade=PrivacyGrade.UNAVAILABLE)