add-or-fix-type-checking

# Add Or Fix Type Checking ## Input - `<target>`: module or directory to type-check (if known). - Optional `make typing` or CI output showing typing failures. ## Workflow 1. **Identify scope from the failing run**: - If you already have `make typing` or CI output, extract the failing file/module paths. - If not, run: ```bash make typing ``` - Choose the narrowest target that covers the failures. 2. **Run `ty check` for the target** to get a focused baseline: ```bash ty check --respect-ignore-files --exclude '**/*_pb*' <target> ``` 3. **Triage errors by category** before fixing anything: - Wrong/missing type annotations on signatures - Attribute access on union types (for example `X | None`) - Functions returning broad unions (for example `str | list | BatchEncoding`) - Mixin/protocol self-type issues - Dynamic attributes on objects or modules - Third-party stub gaps (missing kwargs, missing `__version__`, etc.) 4. **Apply fixes using this priority order** (simplest first): a. **Narrow unions with `isinstance()` / `if x is None` / `hasattr()`**. This is the primary tool for resolving union-type errors. `ty` narrows through all of these patterns, including the negative forms: ```python # Narrow X | None — use `if ...: raise`, never `assert` if x is None: raise ValueError("x must not be None") x.method() # ty knows x is X here # Narrow str | UploadFile if isinstance(field, str): raise TypeError("Expected file upload, got string") await field.read() # ty knows field is UploadFile here # Narrow broad union parameters early in a function body # (common for methods accepting e.g. list | dict | BatchEncoding) if isinstance(encoded_inputs, (list, tuple)): raise TypeError("Expected a mapping, got sequence") encoded_inputs.keys() # ty sees only the dict/mapping types now ``` b. **Use local variables to help ty track narrowing across closures**. When `self.x` is `X | None` and you need to pass it to nested functions or closures, `ty` cannot track that `self.x` stays non-None. Copy to a local variable and narrow the local: ```python manager = self.batching_manager if manager is None: raise RuntimeError("Manager not initialized") # Use `manager` (not `self.batching_manager`) in nested functions ``` c. **Split chained calls when the intermediate type is a broad union**. If `func().method()` fails because `func()` returns a union, split it: ```python # BAD: ty can't narrow through chained calls result = func(return_dict=True).to(device)["input_ids"] # GOOD: split, narrow, then chain result = func(return_dict=True) if not hasattr(result, "to"): raise TypeError("Expected dict-like result") inputs = result.to(device)["input_ids"] ``` d. **Fix incorrect type hints at the source**. If a parameter is typed `X | None` but can never be `None` when actually called, remove `None` from the hint. e. **Annotate untyped attributes**. Add type annotations to instance variables set in `__init__` or elsewhere (for example `self.foo: list[int] = []`). Declare class-level attributes that are set dynamically later (for example `_cache: Cache`, `_token_tensor: torch.Tensor | None`). f. **Use `@overload` for methods with input-dependent return types**. When a method returns different types based on the input type (e.g. `__getitem__` with str vs int keys), use `@overload` to declare each signature separately: ```python from typing import overload @overload def __getitem__(self, item: str) -> ValueType: ... @overload def __getitem__(self, item: int) -> EncodingType: ... @overload def __getitem__(self, item: slice) -> dict[str, ValueType]: ... def __getitem__(self, item: int | str | slice) -> ValueType | EncodingType | dict[str, ValueType]: ... # actual implementation ``` This eliminates `cast()` calls at usage sites by giving the checker precise return types for each call pattern. g. **Make container classes generic to propagate value types**. When a class like `UserDict` holds values whose type changes after transformation (e.g. lists → tensors after `.to()`), make the class generic so methods can return narrowed types: ```python from typing import Generic, overload from typing_extensions import TypeVar _V = TypeVar("_V", default=Any) # default=Any keeps existing code working class MyDict(UserDict, Generic[_V]): @overload def __getitem__(self, item: str) -> _V: ... # ... def to(self, device) -> MyDict[torch.Tensor]: # after .to(), values are tensors ... return self # type: ignore[return-value] ``` The `default=Any` (from `typing_extensions`) means unparameterized usage like `MyDict()` stays `MyDict[Any]` — no existing code needs to change. Only methods that narrow the value type (like `.to()`) declare a specific return type. This eliminates `cast()` at all call sites. h. **Use `self: "ProtocolType"` for mixins**. When a mixin accesses attributes from its host class, define a Protocol in `src/transformers/_typing.py` and annotate `self` on methods that need it. Apply this consistently to all methods in the mixin. Import under `TYPE_CHECKING` to avoid circular imports. i. **Use `TypeGuard` functions for dynamic module attributes** (for example `torch.npu`, `torch.xpu`, `torch.compiler`). Instead of `getattr(torch, "npu")` or `hasattr(torch, "npu") and torch.npu.is_available()`, define a type guard function in `src/transformers/_typing.py`: ```python def has_torch_npu(mod: ModuleType) -> TypeGuard[Any]: return hasattr(mod, "npu") and mod.npu.is_available() ``` Then use it as a narrowing check: `if has_torch_npu(torch): torch.npu.device_count()`. After the guard, `ty` treats the module as `Any`, allowing attribute access without `getattr()` or `cast()`. See existing guards in `_typing.py` for all device backends. **Key rules for type guards**: - Use `TypeGuard[Any]` (not a Protocol) — this is the simplest form that works with `ty` and avoids losing the original module's known attributes. - The guard function must be called directly in an `if` condition for narrowing to work. `ty` does NOT narrow through `and` conditions or `if not guard: return`. - Import guards with `from .._typing import has_torch_xxx` (not via module attribute `_typing.has_torch_xxx`) — `ty` only resolves `TypeGuard` from direct imports. j. **Use `getattr()` / `setattr()` for dynamic model/config attributes**. For runtime-injected fields (for example config/model flags), use `getattr(obj, "field", default)` for reads and `setattr(obj, "field", value)` for writes. Also use `getattr()` for third-party packages missing type stubs (for example `getattr(safetensors, "__version__", "unknown")`). Avoid `getattr(torch, "npu")` style — use type guards instead (see above). k. **Use `cast()` as a last resort before `# type: ignore`**. Use when you've structurally validated the type but the checker can't see it: pattern-matched AST nodes, known-typed dict values, or validated API responses. ```python # After structural validation confirms the type: stmt = cast(cst.Assign, node.body[0]) annotations = cast(list[Annotation], []) ``` Do not use `cast()` for module attribute narrowing — use type guards. Do not use `cast()` when `@overload` or generics can solve it at the source. l. **Use `# type: ignore` only for third-party stub defects**. This means cases where the third-party package's type stubs are wrong or incomplete and there is no way to narrow or cast around it. Examples: - A kwarg that exists at runtime but is missing from the stubs - A method that exists but isn't declared in the stubs Always add the specific error code: `# type: ignore[call-arg]`, not bare `# type: ignore`. 5. **Things to never do**: - **Never use `assert` for type narrowing.** Asserts are stripped by `python -O` and must not be relied on for correctness. Use `if ...: raise` instead. - **Never use `# type: ignore` as a first resort.** Exhaust all approaches above first. - Do not use `getattr(torch, "backend")` to access dynamic device backends (`npu`, `xpu`, `hpu`, `musa`, `mlu`, `neuron`, `compiler`) — use type guards - Do not use `cast()` for module attribute narrowing — use type guards - Do not use `cast()` when `@overload` or generics can eliminate it at the source - Do not add helper methods or abstractions just to satisfy the type checker (especially for only 1-2 occurrences) - Do not pollute base classes with domain-specific fields; use Protocols - Do not add `if x is not None` guards for values guaranteed non-None by the call chain; fix the annotation instead - Do not use conditional inheritance patterns; annotate `self` instead 6. **Organization**: - Keep shared Protocols and type aliases in `src/transformers/_typing.py` - Import type-only symbols under `if TYPE_CHECKING:` to avoid circular deps - Use `from __future__ import annotations` for PEP 604 syntax (`X | Y`) 7. **Verify and close the PR loop**: - Re-run `ty check` on the same `<target>` - Re-run `make typing` to confirm the type/model-rules step passes - If working toward merge readiness, run `make check-repo` - Ensure runtime behavior did not change and run relevant tests 8. **Update CI coverage when adding new typed areas**: - Update `ty_check_dirs` in `Makefile` to include newly type-checked directories.