refactor: tighten checked-source typing coverage

[binaryaaron]

Summary

• Enables the ty missing-override rule for checked source scope.
• Removes internal casts and adds typed helpers across non-PII checked modules.

Test plan

• uv run pytest tests/config/test_parameters.py tests/configurator/test_pydantic_click_options.py tests/data_processing/test_budget.py tests/data_processing/test_data_actions.py tests/data_processing/test_distributions.py tests/preflight/test_plugin_registration.py tests/preflight/test_preflight.py tests/preflight/test_registry_validation.py tests/sdk/test_builder.py -q
• mise run typecheck

Related issue: #614

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• main

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c483210f-2bca-41d4-83a4-0f1989128f6b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch binaryaaron/pr596-04-checked-typing-coverage

---

_{Comment @coderabbitai help to get the list of available commands.}

[greptile-apps]

Greptile Summary

This PR enables the missing-override-decorator ty rule as an error and systematically applies @override across ~35 files covering training, generation, data-processing, evaluation, and observability modules. It also removes most remaining cast calls in favour of typed helpers (TypeIs, TypeGuard, explicit isinstance guards, typed accessor methods), migrates one json.loads + model_validate pair to model_validate_json, and replaces the partial(OpacusDPTrainer, ...) + cast pattern with a proper _TrainerFactory Protocol.

• @override coverage: Added to every subclass method that previously triggered the now-erroring rule, spanning data-action subclasses, trainer callbacks, assemblers, backend classes, and the opacus stubs.
• Cast elimination: cast(RDPAccountant, ...) / cast(PRVAccountant, ...) replaced by guarded accessor methods (rdp_accountant() / prv_accountant()); cast(int, df.columns.get_loc(...)) replaced by an isinstance guard with a ValueError; _TrainerFactory Protocol replaces partial + cast for the DP trainer.
• Type helpers: New TypeIs/TypeGuard narrowing functions (_is_weight_map, is_dataframe, is_iterable, _is_nested_int_list, _is_flat_int_list) remove the need for inline isinstance checks and carry narrowed types forward to callers.

Confidence Score: 5/5

All changes are type-annotation tightening with no logic changes outside the ReplaceDataSource guard and model_validate_json migration, both of which are covered by new tests.

The diff is a mechanical refactoring: @override annotations, cast removal in favour of typed guards and typed accessor methods, and a Protocol-based trainer factory. Each substitution preserves the existing runtime behaviour. The new tests in test_data_actions.py confirm the model_validate_json migration and the ReplaceDataSource column-index guard. The only annotation-level inaccuracy found is the float32 vs float64 distance annotation in the CPU path of kneighbors, which has no runtime impact.

No files require special attention beyond the distance-type annotation in evaluation/nearest_neighbors.py.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/training/huggingface_backend.py	Replaces partial(OpacusDPTrainer, ...) + cast with a typed _TrainerFactory Protocol and factory functions; adds @override throughout — clean, no logic changes.
src/nemo_safe_synthesizer/llm/utils.py	Rewrites auto_map / weight_map processing with match/case and typed helpers (_is_weight_map, AutoMapValue); removes most cast calls; _auto_map_class_refs match is exhaustive given `AutoMapValue = str
src/nemo_safe_synthesizer/generation/vllm_backend.py	Replaces all_equal_type heuristics with explicit TypeGuard helpers (_is_nested_int_list, _is_flat_int_list) and adds @override on key methods; logic is equivalent.
src/nemo_safe_synthesizer/evaluation/nearest_neighbors.py	Narrows return types of fit/kneighbors to npt.NDArray[np.float32] / npt.NDArray[np.int64]; GPU path correctly returns float32 but sklearn CPU path actually returns float64 distances.
src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py	Removes cast(RDPAccountant, ...) in favour of self.accountant.rdp_accountant() accessors; adds @override on all TrainerCallback / DataCollator overrides — safe.
src/nemo_safe_synthesizer/privacy/dp_transformers/privacy_args.py	Adds prv_accountant() / rdp_accountant() accessors with isinstance guards to replace cast calls; accessors are always called under the matching use_prv branch, so guards won't fire in normal operation.
src/nemo_safe_synthesizer/data_processing/actions/data_actions.py	Replaces json.loads + model_validate with model_validate_json; adds isinstance guard on df.columns.get_loc; adds @override on all action subclass methods — correct.
src/nemo_safe_synthesizer/generation/regex_manager.py	Converts dispatch-dict pattern to explicit match/case arms in _type_regex, tightens **kwargs to **kwargs: object throughout — logic unchanged.
src/nemo_safe_synthesizer/utils.py	Adds is_dataframe TypeIs helper, tightens time_function to use ParamSpec, narrows is_iterable to TypeGuard[Iterable[object]], tightens write_json / load_json signatures.
src/nemo_safe_synthesizer/configurator/parameter.py	Fixes __eq__ to return False instead of silently casting NotImplemented; changes comparison return type to `bool
typings/opacus/accountants/init.pyi	Adds @override to RDPAccountant.get_epsilon; minimal stub improvement.
tests/data_processing/test_data_actions.py	New test file covering ReplaceDataSource state JSON round-trip and DatetimeCol state/validate round-trip — good coverage for the model_validate_json migration.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class _TrainerFactory {
        <<Protocol>>
        +__call__(model, processing_class, args, train_dataset, eval_dataset, data_collator, compute_metrics, preprocess_logits_for_metrics, callbacks) Trainer
    }
    class _standard_trainer_factory {
        +__call__(**kwargs) Trainer
    }
    class _opacus_trainer_factory {
        +__call__(privacy_args, true_dataset_size, data_fraction) _TrainerFactory
    }
    class HuggingFaceBackend {
        +trainer_type : _TrainerFactory
        +prepare_params()
        +_build_trainer()
    }
    class SafeSynthesizerAccountant {
        +accountant : PRVAccountant | RDPAccountant
        +prv_accountant() PRVAccountant
        +rdp_accountant() RDPAccountant
        +compute_epsilon() float
    }
    _TrainerFactory <|.. _standard_trainer_factory : implements
    _opacus_trainer_factory --> _TrainerFactory : returns
    HuggingFaceBackend --> _TrainerFactory : uses
    HuggingFaceBackend --> _standard_trainer_factory : default
    HuggingFaceBackend --> _opacus_trainer_factory : when DP enabled

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
classDiagram
    class _TrainerFactory {
        <<Protocol>>
        +__call__(model, processing_class, args, train_dataset, eval_dataset, data_collator, compute_metrics, preprocess_logits_for_metrics, callbacks) Trainer
    }
    class _standard_trainer_factory {
        +__call__(**kwargs) Trainer
    }
    class _opacus_trainer_factory {
        +__call__(privacy_args, true_dataset_size, data_fraction) _TrainerFactory
    }
    class HuggingFaceBackend {
        +trainer_type : _TrainerFactory
        +prepare_params()
        +_build_trainer()
    }
    class SafeSynthesizerAccountant {
        +accountant : PRVAccountant | RDPAccountant
        +prv_accountant() PRVAccountant
        +rdp_accountant() RDPAccountant
        +compute_epsilon() float
    }
    _TrainerFactory <|.. _standard_trainer_factory : implements
    _opacus_trainer_factory --> _TrainerFactory : returns
    HuggingFaceBackend --> _TrainerFactory : uses
    HuggingFaceBackend --> _standard_trainer_factory : default
    HuggingFaceBackend --> _opacus_trainer_factory : when DP enabled

Loading

_{Reviews (2): Last reviewed commit: "refactor(training): narrow trainer facto..." | Re-trigger Greptile}

[greptile-apps]

Missing implicit fallthrough in _auto_map_class_refs

The old implementation had an explicit return [] for any non-str/non-list value. The refactored match statement covers str() and list() — which are exhaustive over AutoMapValue = str | list[object] — but has no case _: arm. If the function is ever called with a value that slips past the type alias (e.g., from untyped dynamic code), Python silently returns None instead of [], violating the declared -> list[str] return type. Adding case _: return [] restores the original defensive behaviour at zero cost.