refactor: improve type-safe config and NER boundaries

.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"
 

pyproject.toml

@@ -361,15 +361,15 @@ safe-synthesizer = "nemo_safe_synthesizer.cli.cli:cli"
 # Platform-specific imports (e.g. vLLM on macOS) use ty: ignore[unresolved-import]
 # that appear unused when the package is installed. Suppress the noise.
 unused-ignore-comment = "ignore"
+missing-override-decorator = "error"
+
 
 [tool.ty.environment]
 extra-paths = ["typings"]
 
 [tool.ty.src]
 # Below is a list of excluded directories from ty typechecks.
 exclude = [
-    ".uv_cache", # Cache Dir in CI
-    "./docs/**/*.ipynb",
-    "./src/nemo_safe_synthesizer/pii_replacer/",
-    "./tools/*.py"
+  "./uv-cache",
+  "./docs/**/*.ipynb",
 ]

script/slurm/nss_top.py

@@ -29,6 +29,7 @@
 from textual.containers import Vertical
 from textual.reactive import reactive
 from textual.widgets import DataTable, Footer, Header, RichLog, Static
+from typing_extensions import override
 
 # squeue format fields and matching column headers
 _SQUEUE_FIELDS = ["%i", "%j", "%T", "%P", "%M", "%l", "%D", "%C", "%R"]
@@ -176,6 +177,7 @@ def __init__(self, user: str, log_dir: Path | None, refresh_secs: int) -> None:
 
     # ── Layout ────────────────────────────────────────────────────────────────
 
+    @override
     def compose(self) -> ComposeResult:
         yield Header(show_clock=True)
         with Vertical():

src/nemo_safe_synthesizer/cli/artifact_structure.py

@@ -27,6 +27,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Generic, Self, TypeVar, overload
 
+from typing_extensions import override
+
 from ..observability import get_logger
 from ..utils import write_json
 
@@ -249,17 +251,21 @@ def path(self) -> Path:
         """The resolved directory path."""
         return self._path
 
+    @override
     def __fspath__(self) -> str:
         """Support os.fspath() for use with open(), etc."""
         return str(self._path)
 
+    @override
     def __str__(self) -> str:
         """Return the path as a string."""
         return str(self._path)
 
+    @override
     def __repr__(self) -> str:
         return f"BoundDir({self._path!r})"
 
+    @override
     def __eq__(self, other: object) -> bool:
         """Support comparison with Path objects."""
         if isinstance(other, BoundDir):
@@ -268,6 +274,7 @@ def __eq__(self, other: object) -> bool:
             return self._path == other
         return NotImplemented
 
+    @override
     def __hash__(self) -> int:
         return hash(self._path)
 

src/nemo_safe_synthesizer/cli/datasets.py

@@ -21,6 +21,16 @@
 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:
+    return value
+
+
 class DatasetInfo(BaseModel):
     """Entry in the dataset registry."""
 
@@ -96,20 +106,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 +130,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/cli/run.py

@@ -9,7 +9,7 @@
 import sys
 from collections.abc import Callable
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
 
 import click
 
@@ -35,8 +35,11 @@
     from ..sdk.library_builder import SafeSynthesizer
     from .artifact_structure import Workdir
 
+P = ParamSpec("P")
+R = TypeVar("R")
 
-def common_run_options(f: Callable[..., object]) -> Callable[..., object]:
+
+def common_run_options(f: Callable[P, R]) -> Callable[P, R]:
     """Decorator to add common options for run commands.
 
     Apply this above ``@pydantic_options`` in source order. Python applies

src/nemo_safe_synthesizer/cli/utils.py

@@ -24,6 +24,7 @@
 from pydantic import ValidationError
 
 from ..config import SafeSynthesizerParameters
+from ..config.parameters import ConfigPatch
 from ..defaults import DEFAULT_ARTIFACTS_PATH
 from ..observability import configure_logging_from_workdir, get_logger, initialize_observability
 from ..utils import merge_dicts
@@ -431,7 +432,7 @@ def _initialize_logging_for_cli_from_settings(
     return run_logger
 
 
-def merge_overrides(config_path: str | Path | None, overrides: dict) -> SafeSynthesizerParameters:
+def merge_overrides(config_path: str | Path | None, overrides: ConfigPatch) -> SafeSynthesizerParameters:
     """Merge overrides into a SafeSynthesizerParameters object.
 
     If config_path is None, use the overrides to create a new SafeSynthesizerParameters object.
@@ -446,11 +447,9 @@ def merge_overrides(config_path: str | Path | None, overrides: dict) -> SafeSynt
     """
     try:
         if config_path is None:
-            my_config = SafeSynthesizerParameters.model_validate(overrides)
+            my_config = SafeSynthesizerParameters.from_config_patch(overrides)
         else:
-            file_config = SafeSynthesizerParameters.from_yaml(config_path).model_dump(exclude_unset=True)
-            params = merge_dicts(file_config, overrides)
-            my_config = SafeSynthesizerParameters.model_validate(params)
+            my_config = SafeSynthesizerParameters.from_yaml(config_path).with_config_patch(overrides)
     except ValidationError as e:
         click.echo(f"{config_path} is invalid:\n{e}")
         sys.exit(1)

src/nemo_safe_synthesizer/config/autoconfig.py

@@ -20,8 +20,7 @@
 from ..defaults import DEFAULT_MAX_SEQ_LENGTH, MAX_ROPE_SCALING_FACTOR
 from ..llm.metadata import ModelMetadata
 from ..observability import get_logger
-from ..utils import merge_dicts
-from .parameters import SafeSynthesizerParameters
+from .parameters import ConfigPatch, SafeSynthesizerParameters
 from .types import AUTO_STR
 
 if TYPE_CHECKING:
@@ -304,14 +303,13 @@ def _build_updated_params(
         Returns:
             The validated SafeSynthesizerParameters.
         """
-        new_params = {
+        new_params: ConfigPatch = {
             "training": training_params,
             "data": data_params,
             "privacy": privacy_params,
         }
-        updated_params = merge_dicts(self._config.model_dump(exclude_unset=True), new_params)
-        logger.debug(f"params to update: {updated_params}")
-        my_config = SafeSynthesizerParameters.model_validate(updated_params)
+        logger.debug(f"params to update: {new_params}")
+        my_config = self._config.with_config_patch(new_params)
         logger.debug(f"auto-updated config: {my_config.model_dump(exclude_unset=True)}")
         return my_config
 

src/nemo_safe_synthesizer/config/base.py

@@ -9,6 +9,7 @@
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict
+from typing_extensions import override
 
 __all__ = ["NSSBaseModel", "LRScheduler", "pydantic_model_config"]
 
@@ -34,6 +35,7 @@ class NSSBaseModel(BaseModel):
 
     model_config = pydantic_model_config
 
+    @override
     def dict(self, **kwargs: Any) -> dict[str, Any]:  # ty: ignore[invalid-type-form] -- method name shadows builtin dict; backward-compat shim for model_dump
         """Return a dict representation via ``model_dump`` for backward compatibility."""
         return self.model_dump(**kwargs)

src/nemo_safe_synthesizer/config/parameters.py

@@ -4,11 +4,14 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Self
+from collections.abc import Mapping
+from typing import Any, Self, TypeAlias
 
 from pydantic import BaseModel, Field, model_validator
+from typing_extensions import override
 
 from ..configurator.parameters import Parameters
+from ..data_processing.records.json_types import JsonValue
 from ..errors import ParameterError
 from ..observability import get_logger
 from ..telemetry import _telemetry_enabled
@@ -23,13 +26,16 @@
 from .training import TrainingHyperparams
 from .types import AUTO_STR
 
-__all__ = ["SafeSynthesizerParameters"]
+ConfigPatch: TypeAlias = Mapping[str, JsonValue]
+_SectionPatch: TypeAlias = dict[str, object]
+
+__all__ = ["ConfigPatch", "SafeSynthesizerParameters"]
 
 
 logger = get_logger(__name__)
 
 
-def _collect_set_fields(model: BaseModel) -> dict[str, Any]:
+def _collect_set_fields(model: BaseModel) -> _SectionPatch:
     """Recursively collect a model's explicitly-set fields as a nested dict.
 
     Unlike ``model_dump(exclude_unset=True)``, nested models are always
@@ -38,7 +44,7 @@ def _collect_set_fields(model: BaseModel) -> dict[str, Any]:
     ``cfg.generation.validation.foo = True``) are captured. A nested model is
     included only when it has at least one set field of its own.
     """
-    overrides: dict[str, Any] = {}
+    overrides: _SectionPatch = {}
     for name in type(model).model_fields:
         value = getattr(model, name)
         if isinstance(value, BaseModel):
@@ -192,6 +198,7 @@ def check_timeseries_group_column(self) -> Self:
         return self
 
     @classmethod
+    @override
     def from_params(cls, **kwargs) -> "SafeSynthesizerParameters":
         """Convert singular, flat parameters to nested structure.
 
@@ -237,6 +244,21 @@ def from_params(cls, **kwargs) -> "SafeSynthesizerParameters":
             extra["emit_telemetry"] = kwargs["emit_telemetry"]
         return cls(**extra)
 
+    @classmethod
+    def from_config_patch(cls, patch: ConfigPatch) -> Self:
+        """Validate a sparse top-level config patch as a full configuration."""
+        return cls.model_validate(patch)
+
+    def with_config_patch(self, patch: ConfigPatch) -> Self:
+        """Apply a sparse top-level config patch and revalidate the result.
+
+        Only fields explicitly set on ``self`` are carried into the merge before
+        applying ``patch``. This preserves file/CLI precedence while keeping
+        default values implicit for future ``exclude_unset`` dumps.
+        """
+        params = merge_dicts(self.model_dump(exclude_unset=True), patch)
+        return type(self).model_validate(params)
+
     def with_runtime_overrides(self, runtime: SafeSynthesizerParameters) -> "SafeSynthesizerParameters":
         """Apply resume-time generation/evaluation/telemetry overrides onto a copy of self.
 

src/nemo_safe_synthesizer/configurator/parameter.py

@@ -13,10 +13,12 @@
 import operator
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import Any, Generic, TypeVar, cast, get_args
+from types import NotImplementedType
+from typing import Any, Generic, TypeVar, get_args
 
 from pydantic import BaseModel, GetCoreSchemaHandler, model_serializer
 from pydantic_core import core_schema
+from typing_extensions import override
 
 DataT = TypeVar(
     "DataT", bound=(int | float | str | bytes | bool | None | Sequence[int | float | str | bytes | bool | BaseModel])
@@ -55,6 +57,7 @@ def ser_model(self) -> "dict[str, DataT] | DataT | Sequence[DataT] | Parameter[D
         else:
             return self
 
+    @override
     def __str__(self):
         return self.__repr__()
 
@@ -86,7 +89,7 @@ def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler
         non_instance_schema = core_schema.no_info_before_validator_function(cls, sequence_t_schema)
         return core_schema.union_schema([instance_schema, non_instance_schema])
 
-    def _comp_helper(self, other: "Parameter[DataT] | DataT", op: Callable[[Any, Any], bool]) -> bool | None:
+    def _comp_helper(self, other: object, op: Callable[[Any, Any], bool]) -> bool | NotImplementedType:
         """Apply a comparison operator ``op`` to ``self.value`` and the unwrapped value of ``other``."""
         match other:
             case Parameter(value=y) if isinstance(self.value, type(y)):
@@ -108,6 +111,7 @@ def __gt__(self, other: "Parameter[DataT] | DataT") -> bool | None:
     def __lt__(self, other: "Parameter[DataT] | DataT") -> bool | None:
         return self._comp_helper(other, operator.__lt__)
 
+    @override
     def __eq__(self, other: object) -> bool:
-        result = self._comp_helper(cast("Parameter[DataT] | DataT", other), operator.__eq__)
-        return cast(bool, result)
+        result = self._comp_helper(other, operator.__eq__)
+        return False if result is NotImplemented else result