type: clean mypy errors in inspect, loader, onnx, optim and enforce in CI

.github/workflows/lint.yml

@@ -13,7 +13,7 @@ concurrency:
 jobs:
   lint:
     runs-on: windows-latest
-    # Bumped from 5: combined mypy on 12 packages cold-starts at ~3-4 min on
+    # Bumped from 5: combined mypy on 16 packages cold-starts at ~3-4 min on
     # Windows runners; the original 5-min ceiling cancelled mid-run.
     timeout-minutes: 10
 
@@ -60,3 +60,7 @@ jobs:
           -p winml.modelkit.eval
           -p winml.modelkit.export
           -p winml.modelkit.inference
+          -p winml.modelkit.inspect
+          -p winml.modelkit.loader
+          -p winml.modelkit.onnx
+          -p winml.modelkit.optim

src/winml/modelkit/inspect/formatter.py

@@ -17,7 +17,7 @@
 from rich.table import Table
 from rich.tree import Tree
 
-from .types import InspectResult, ModuleInfo, SupportLevel
+from .types import HierarchyInfo, InspectResult, ModuleInfo, SupportLevel
 
 
 # Status icons
@@ -505,7 +505,7 @@ def output_json(result: InspectResult, verbose: bool = False) -> str:
     return json.dumps(data, indent=2)
 
 
-def _hierarchy_to_dict(hierarchy) -> dict[str, Any]:
+def _hierarchy_to_dict(hierarchy: HierarchyInfo) -> dict[str, Any]:
     """Convert HierarchyInfo to dictionary.
 
     Args:

src/winml/modelkit/inspect/hierarchy.py

@@ -174,8 +174,13 @@ def process_module(
         module: nn.Module,
         depth: int,
         parent_path: str = "",
-    ) -> ModuleInfo | None:
-        """Recursively process modules, filtering out torch.nn modules."""
+    ) -> ModuleInfo | list[ModuleInfo] | None:
+        """Recursively process modules, filtering out torch.nn modules.
+
+        Returns a single ``ModuleInfo`` for an HF module, a ``list[ModuleInfo]``
+        of HF descendants propagated up through a skipped ``nn`` module, or
+        ``None`` when no HF module is found in this subtree.
+        """
         nonlocal hf_module_count, nn_module_count
 
         full_path = f"{parent_path}.{name}" if parent_path else name

src/winml/modelkit/inspect/module_io_capture.py

@@ -13,12 +13,16 @@
 import inspect
 import logging
 from dataclasses import dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from torch import nn
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -105,9 +109,9 @@ def capture_module_io(
 
         def make_hook(
             mod_name: str, cls_name: str, fwd_params: list[str],
-        ):
+        ) -> Callable[[nn.Module, tuple[Any, ...], dict[str, Any], Any], None]:
             def hook(
-                mod: nn.Module, args: tuple, kwargs: dict, output: Any,
+                mod: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any,
             ) -> None:
                 info = ModuleIOInfo(class_name=cls_name, module_path=mod_name)
 

src/winml/modelkit/inspect/resolver.py

@@ -10,7 +10,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 from ..loader.resolution import _get_custom_model_class
 from ..loader.task import (
@@ -269,7 +269,14 @@ def resolve_exporter(
     # Check MODEL_BUILD_CONFIGS for predefined config
     if model_type_normalized in MODEL_BUILD_CONFIGS:
         config: WinMLBuildConfig = MODEL_BUILD_CONFIGS[model_type_normalized]
+        # MODEL_BUILD_CONFIGS entries are HF export configs; export is None only on
+        # the direct-ONNX build path, which never reaches this registry lookup.
         export_config = config.export
+        if export_config is None:
+            raise ValueError(
+                f"MODEL_BUILD_CONFIGS entry for {model_type_normalized!r} is missing an "
+                "export config (export is None only on the direct-ONNX build path)."
+            )
 
         # Extract input tensors
         input_tensors: list[TensorInfo] = []
@@ -326,8 +333,8 @@ def resolve_exporter(
                 config_name = onnx_config_cls.__name__
 
             # Extract tensor specs via resolve_io_specs (shared with config command)
-            input_tensors: list[TensorInfo] = []
-            output_tensors: list[TensorInfo] = []
+            input_tensors = []
+            output_tensors = []
 
             if hf_config is not None:
                 try:
@@ -695,8 +702,12 @@ def resolve_io_config(
 
     def get_config_attr(
         attr_name: str,
-    ) -> int | tuple[int, int] | list | None:
-        """Get attribute from main config or any nested config."""
+    ) -> Any:
+        """Get attribute from main config or any nested config.
+
+        Returns ``Any``: HF config attributes are dynamically typed (int, tuple,
+        list, str, ...), so each call site narrows to its target field type.
+        """
         value = getattr(config, attr_name, None)
         if value is not None:
             return value
@@ -952,7 +963,7 @@ def _resolve_processor_from_hub_configs(model_id: str) -> _HubConfigResult:
     from pathlib import Path
 
     from huggingface_hub import hf_hub_download
-    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
+    from huggingface_hub.errors import EntryNotFoundError, RepositoryNotFoundError
 
     processor_class: str | None = None
     tokenizer_class: str | None = None

src/winml/modelkit/loader/__init__.py

@@ -25,6 +25,8 @@
     See docs/design/loader/hf.md for design details.
 """
 
+from typing import Any
+
 from .config import WinMLLoaderConfig, resolve_loader_config
 from .resolution import TaskResolution, TaskSource, resolve_composite, resolve_task
 from .task import (
@@ -65,7 +67,7 @@
 }
 
 
-def __getattr__(name: str):
+def __getattr__(name: str) -> Any:
     """Lazy-load heavy exports (hf.py imports transformers)."""
     if name in _LAZY_IMPORTS:
         module_path, attr_name = _LAZY_IMPORTS[name]

src/winml/modelkit/loader/config.py

@@ -13,7 +13,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 
 if TYPE_CHECKING:
@@ -266,7 +266,7 @@ def _create_hf_config_from_model_class(model_class: type) -> PretrainedConfig:
         )
     hf_config = config_cls()
     hf_config.architectures = [model_class.__name__]
-    return hf_config
+    return cast("PretrainedConfig", hf_config)
 
 
 def _resolve_hf_config_for_class(

src/winml/modelkit/loader/hf.py

@@ -27,7 +27,7 @@
 import importlib.util
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast
 
 from transformers import AutoConfig
 
@@ -76,7 +76,7 @@ def resolve_hf_model_class(class_name: str) -> type:
         cls = getattr(module, class_name, None)
         if cls is not None:
             logger.debug("Resolved '%s' from '%s'", class_name, module_name)
-            return cls
+            return cast("type", cls)
 
     raise ImportError(
         f"Model class '{class_name}' not found in any of: {', '.join(_HF_MODEL_MODULES)}"
@@ -222,6 +222,9 @@ def load_hf_model(
     from .resolution import resolve_task
 
     if user_script is not None:
+        # model_class is guaranteed non-None here: the validation block above
+        # raises when user_script is set without a model_class.
+        assert model_class is not None
         resolved_class = _load_class_from_script(user_script, model_class)
         logger.info("Using custom model class from script: %s", model_class)
         # Surfaced modality-aware task (consistent with the non-script branch).
@@ -237,7 +240,11 @@ def load_hf_model(
 
     # [4] Model Instantiation
     logger.debug("Loading model with class: %s", resolved_class.__name__)
-    model = resolved_class.from_pretrained(
+    # resolved_class is a dynamically-resolved model class (transformers, timm,
+    # diffusers, ...); from_pretrained is a duck-typed boundary across these libs,
+    # so go through an Any-typed alias rather than a static attribute access.
+    loader_cls: Any = resolved_class
+    model = loader_cls.from_pretrained(
         model_name_or_path,
         trust_remote_code=trust_remote_code,
     )

src/winml/modelkit/loader/resolution.py

@@ -20,7 +20,7 @@
 import logging
 from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from .task import (
     HF_TASK_DEFAULTS,
@@ -123,7 +123,7 @@ class from the ``transformers`` package.
 
     try:
         transformers_module = importlib.import_module("transformers")
-        return getattr(transformers_module, arch_name)
+        return cast("type", getattr(transformers_module, arch_name))
     except AttributeError as e:
         raise ValueError(
             f"Cannot import {arch_name} from transformers. Please specify task explicitly."
@@ -145,7 +145,7 @@ def _detect_task_from_model_class(model_class: type) -> str:
     """
     from optimum.exporters.tasks import TasksManager
 
-    return TasksManager.infer_task_from_model(model_class)
+    return cast("str", TasksManager.infer_task_from_model(model_class))
 
 
 def _upgrade_fill_mask_for_seq2seq(task: str, config: PretrainedConfig) -> str:
@@ -196,6 +196,8 @@ def _resolve_task_modality(config: PretrainedConfig, task: str) -> str:
     except ValueError:
         return task
     main_input = getattr(model_class, "main_input_name", None)
+    if main_input is None:
+        return task
     return _FEATURE_MODALITY_BY_MAIN_INPUT.get(main_input, task)
 
 
@@ -230,7 +232,7 @@ def _get_custom_model_class(model_type: str, task: str) -> type | None:
     if task in HF_TASK_DEFAULTS:
         import transformers
 
-        return getattr(transformers, HF_TASK_DEFAULTS[task])
+        return cast("type", getattr(transformers, HF_TASK_DEFAULTS[task]))
 
     return None
 
@@ -316,7 +318,7 @@ def _composite_components_for_task(model_type: str, task: str) -> CompositeCompo
     from ..models.winml import WinMLCompositeModel
     from ..models.winml.composite_model import COMPOSITE_MODEL_REGISTRY
 
-    distinct: dict[tuple, type] = {}
+    distinct: dict[tuple, type[WinMLCompositeModel]] = {}
     for (m_type, reg_task), cls in COMPOSITE_MODEL_REGISTRY.items():
         if m_type != model_type or not issubclass(cls, WinMLCompositeModel):
             continue
@@ -350,6 +352,10 @@ def resolve_task(
     model_type_norm = model_type.lower().replace("_", "-") if model_type else ""
     model_id = getattr(config, "_name_or_path", "") or None
 
+    # Declared once up front so the Stage-0 branches (which assign a concrete str)
+    # and the Stage-1 detection (which starts at None) share one str | None type.
+    opt_task: str | None = None
+
     # --- Stage 0: user override (short-circuits detection) ----------------
     if model_class is not None:
         if task is not None:
@@ -402,7 +408,7 @@ def resolve_task(
         )
 
     # --- Stage 1: detection -----------------------------------------------
-    opt_task: str | None = None
+    # opt_task stays at its hoisted None until a detection sub-stage sets it.
     source: TaskSource | None = None
     resolved = None
 

src/winml/modelkit/loader/task.py

@@ -20,6 +20,11 @@
 
 from __future__ import annotations
 
+import logging
+from typing import cast
+
+
+logger = logging.getLogger(__name__)
 
 # Task abbreviations for cache keys (47 tasks from HuggingFace Transformers)
 TASK_ABBREV: dict[str, str] = {
@@ -214,7 +219,7 @@ def normalize_task(task: str) -> str:
     """
     from optimum.exporters.tasks import TasksManager
 
-    return TasksManager.map_from_synonym(task)
+    return cast("str", TasksManager.map_from_synonym(task))
 
 
 # WinML task-synonym extensions — extend Optimum's ``TasksManager.map_from_synonym``
@@ -255,7 +260,7 @@ def to_optimum_task(task: str) -> str:
 
     from optimum.exporters.tasks import TasksManager
 
-    return TasksManager.map_from_synonym(task)
+    return cast("str", TasksManager.map_from_synonym(task))
 
 
 def get_task_abbrev(task: str) -> str:

src/winml/modelkit/onnx/__init__.py

@@ -13,6 +13,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from .domains import ONNXDomain
 from .dtypes import SupportedONNXType, remove_optional_from_type_annotation
 from .external_data import copy_onnx_model
@@ -53,7 +55,7 @@
 }
 
 
-def __getattr__(name: str):
+def __getattr__(name: str) -> Any:
     """Lazy-load detection module to avoid circular import with compiler."""
     if name in _LAZY_IMPORTS:
         module_path, attr_name = _LAZY_IMPORTS[name]

src/winml/modelkit/onnx/domains.py

@@ -57,7 +57,7 @@ def _init_custom_schemas(domains: set[str] | None = None) -> None:
         return
 
     try:
-        from onnxruntime.capi.onnxruntime_pybind11_state import (  # type: ignore[import]
+        from onnxruntime.capi.onnxruntime_pybind11_state import (
             get_all_operator_schema,
         )
     except ImportError:
@@ -85,7 +85,8 @@ def _init_custom_schemas(domains: set[str] | None = None) -> None:
                     inp.name,
                     inp.typeStr,
                     inp.description,
-                    param_option=OpSchema.FormalParameterOption(inp.option.value),
+                    # onnx's pybind11 enum stubs don't model value-lookup construction
+                    param_option=OpSchema.FormalParameterOption(inp.option.value),  # type: ignore[call-arg]
                 )
                 for inp in ort_schema.inputs
             ]
@@ -94,7 +95,8 @@ def _init_custom_schemas(domains: set[str] | None = None) -> None:
                     out.name,
                     out.typeStr,
                     out.description,
-                    param_option=OpSchema.FormalParameterOption(out.option.value),
+                    # onnx's pybind11 enum stubs don't model value-lookup construction
+                    param_option=OpSchema.FormalParameterOption(out.option.value),  # type: ignore[call-arg]
                 )
                 for out in ort_schema.outputs
             ]
@@ -105,7 +107,8 @@ def _init_custom_schemas(domains: set[str] | None = None) -> None:
             attributes = [
                 OpSchema.Attribute(
                     attr.name,
-                    OpSchema.AttrType(attr.type.value),
+                    # onnx's pybind11 enum stubs don't model value-lookup construction
+                    OpSchema.AttrType(attr.type.value),  # type: ignore[call-arg]
                     attr.description,
                     required=attr.required,
                 )