add-or-fix-type-checking
Add Or Fix Type Checking
Input
<target>: module or directory to type-check (if known).- Optional
make typingor CI output showing typing failures.
Workflow
-
Identify scope from the failing run:
- If you already have
make typingor CI output, extract the failing file/module paths. - If not, run:
make typing - Choose the narrowest target that covers the failures.
- If you already have
-
Run
ty checkfor the target to get a focused baseline:ty check --respect-ignore-files --exclude '**/*_pb*' <target> -
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.)
-
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.tynarrows through all of these patterns, including the negative forms:# 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 nowb. Use local variables to help ty track narrowing across closures. When
self.xisX | Noneand you need to pass it to nested functions or closures,tycannot track thatself.xstays non-None. Copy to a local variable and narrow the local:manager = self.batching_manager if manager is None: raise RuntimeError("Manager not initialized") # Use `manager` (not `self.batching_manager`) in nested functionsc. Split chained calls when the intermediate type is a broad union. If
func().method()fails becausefunc()returns a union, split it:# 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 | Nonebut can never beNonewhen actually called, removeNonefrom the hint.e. Annotate untyped attributes. Add type annotations to instance variables set in
__init__or elsewhere (for exampleself.foo: list[int] = []). Declare class-level attributes that are set dynamically later (for example_cache: Cache,_token_tensor: torch.Tensor | None).f. Use
@overloadfor 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@overloadto declare each signature separately: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 implementationThis 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
UserDictholds values whose type changes after transformation (e.g. lists → tensors after.to()), make the class generic so methods can return narrowed types: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(fromtyping_extensions) means unparameterized usage likeMyDict()staysMyDict[Any]— no existing code needs to change. Only methods that narrow the value type (like.to()) declare a specific return type. This eliminatescast()at all call sites.h. Use
self: "ProtocolType"for mixins. When a mixin accesses attributes from its host class, define a Protocol insrc/transformers/_typing.pyand annotateselfon methods that need it. Apply this consistently to all methods in the mixin. Import underTYPE_CHECKINGto avoid circular imports.i. Use
TypeGuardfunctions for dynamic module attributes (for exampletorch.npu,torch.xpu,torch.compiler). Instead ofgetattr(torch, "npu")orhasattr(torch, "npu") and torch.npu.is_available(), define a type guard function insrc/transformers/_typing.py: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,tytreats the module asAny, allowing attribute access withoutgetattr()orcast(). See existing guards in_typing.pyfor all device backends.Key rules for type guards:
- Use
TypeGuard[Any](not a Protocol) — this is the simplest form that works withtyand avoids losing the original module's known attributes. - The guard function must be called directly in an
ifcondition for narrowing to work.tydoes NOT narrow throughandconditions orif not guard: return. - Import guards with
from .._typing import has_torch_xxx(not via module attribute_typing.has_torch_xxx) —tyonly resolvesTypeGuardfrom direct imports.
j. Use
getattr()/setattr()for dynamic model/config attributes. For runtime-injected fields (for example config/model flags), usegetattr(obj, "field", default)for reads andsetattr(obj, "field", value)for writes. Also usegetattr()for third-party packages missing type stubs (for examplegetattr(safetensors, "__version__", "unknown")). Avoidgetattr(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.# 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 usecast()when@overloador generics can solve it at the source.l. Use
# type: ignoreonly 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.
- Use
-
Things to never do:
- Never use
assertfor type narrowing. Asserts are stripped bypython -Oand must not be relied on for correctness. Useif ...: raiseinstead. - Never use
# type: ignoreas 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@overloador 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 Noneguards for values guaranteed non-None by the call chain; fix the annotation instead - Do not use conditional inheritance patterns; annotate
selfinstead
- Never use
-
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 annotationsfor PEP 604 syntax (X | Y)
- Keep shared Protocols and type aliases in
-
Verify and close the PR loop:
- Re-run
ty checkon the same<target> - Re-run
make typingto 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
- Re-run
-
Update CI coverage when adding new typed areas:
- Update
ty_check_dirsinMakefileto include newly type-checked directories.
- Update