diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index ce9eb391d558..d3aa2e846a45 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -24,7 +24,23 @@ `MappingProxyType(...)`) are not construction and pass, as does the value passed directly to a wrapper: it is frozen before it can escape, though anything mutable nested inside it still counts. Annotation-internal lists - (`Callable[[int], str]`) are exempt. Suppress with `# mutable-ok: `. + (`Callable[[int], str]`) are exempt, as is a dict display assigned directly + to a name annotated with a same-module all-ReadOnly TypedDict + (`x: Final[MyTd] = {...}`, bare or under `Final[...]`): it is the literal + spelling of the `MyTd(...)` call, which was never construction, and an + all-ReadOnly payload cannot be grown or rewritten. The exemption is + conservative: the TypedDict's name must be bound exactly once in the file + (a name also bound in any other form, another def, an assignment or loop + target, an import alias, a parameter, might resolve to something mutable + at the annotation site, and a `from x import *` anywhere in the file + disqualifies every name, since what it binds is statically invisible), + every field it declares or inherits in-module must be + `ReadOnly[...]`, only the display itself is exempt (nested mutables still + count), a TypedDict imported from another module or nested inside a def + or class is out of reach (only ones defined at the module's top level + resolve file-wide), and a dotted annotation (`x: mod.Td = {...}`) never + matches, since it cannot name a local class. Suppress with + `# mutable-ok: `. LIT003 noqa suppression without rule codes or without a reason. Required shape: `# noqa: TID251 # ` LIT004 pyright/mypy ignore without bracketed codes or without a reason. @@ -485,6 +501,92 @@ def _frozen_argument_ids(tree: ast.AST) -> frozenset[int]: ) +def _declared_type_name(annotation: ast.expr) -> str | None: + """The bare name an AnnAssign annotation spells, looking through a `Final[...]` wrapper. + + Qualified annotations (`x: mod.Td`) yield None: a dotted name refers to another + module's attribute, which can never be a class defined in the file being checked, + so reducing it to its tail would let an imported type borrow a local one's name. + """ + if isinstance(annotation, ast.Subscript) and _head_name(annotation.value) == "Final": + return annotation.slice.id if isinstance(annotation.slice, ast.Name) else None + return annotation.id if isinstance(annotation, ast.Name) else None + + +def _binding_names(node: ast.AST) -> tuple[str, ...]: + """The names one node binds: def/class/import/parameter/global/except/match forms, + plus any Name stored or deleted, which covers every assignment, loop, walrus, + unpacking, comprehension, and `with` target.""" + if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + return (node.name,) + if isinstance(node, (ast.Import, ast.ImportFrom)): + return tuple(alias.asname or alias.name.partition(".")[0] for alias in node.names) + if isinstance(node, (ast.Global, ast.Nonlocal)): + return tuple(node.names) + if isinstance(node, ast.arg): + return (node.arg,) + if isinstance(node, ast.ExceptHandler) and node.name is not None: + return (node.name,) + if isinstance(node, (ast.MatchAs, ast.MatchStar)) and node.name is not None: + return (node.name,) + if isinstance(node, ast.MatchMapping) and node.rest is not None: + return (node.rest,) + if isinstance(node, ast.Name) and isinstance(node.ctx, (ast.Store, ast.Del)): + return (node.id,) + return () + + +def _typeddict_assigned_value_ids(tree: ast.AST) -> frozenset[int]: + """ids() of every dict display assigned directly to a frozen-TypedDict-annotated name. + + `x: MyTd = {...}` is the literal spelling of the `MyTd(...)` call, which was + never construction to begin with, so the display is a one-shot build -- provided + the annotation provably names a frozen payload. Three conditions gate that: + MyTd is a class-form TypedDict at the module's top level (imported ones are + invisible, exactly as in LIT012's base-class resolution, and one nested in a + def or class is skipped, since its name does not resolve outside the scope + that defines it, while a top-level one resolves everywhere); its name is + bound exactly once + in the file (resolution here is scope-blind, so a name the file also binds in + any other form, another def, an assignment or loop target, an import alias, a + parameter, could resolve to something mutable at the annotation site, and a + `from x import *` anywhere disqualifies every name in the file, since what it + binds is statically invisible); and every field it declares or inherits + within the module is `ReadOnly[...]`, so no holder can statically + rewrite a key even where LIT012 was suppressed or is riding its budget. Only + the display itself is exempt; anything mutable nested inside it still trips + LIT002. + """ + top_level_ids = frozenset(id(stmt) for stmt in (tree.body if isinstance(tree, ast.Module) else ())) + classes = tuple(cls for cls in _typeddict_classes(tree) if id(cls) in top_level_ids) + by_name = {cls.name: cls for cls in classes} + + def frozen_lineage(name: str, seen: frozenset[str]) -> bool: + cls = by_name.get(name) + if cls is None or name in seen: + return False + if not all(_has_readonly_qualifier(field.annotation) for field in _class_fields(cls)): + return False + return all( + base == TYPEDDICT_BASE or frozen_lineage(base, seen | frozenset((name,))) + for base in _base_names(cls) + ) + + bindings = tuple(name for node in ast.walk(tree) for name in _binding_names(node)) + if "*" in bindings: + return frozenset() + frozen_names = frozenset( + cls.name for cls in classes if bindings.count(cls.name) == 1 and frozen_lineage(cls.name, frozenset()) + ) + return frozenset( + id(node.value) + for node in ast.walk(tree) + if isinstance(node, ast.AnnAssign) + and isinstance(node.value, ast.Dict) + and _declared_type_name(node.annotation) in frozen_names + ) + + def _construction_kind(node: ast.expr) -> str | None: """Human label if `node` builds a mutable collection, else None.""" if isinstance(node, ast.List): @@ -509,10 +611,9 @@ def _construction_kind(node: ast.expr) -> str | None: def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: - in_annotation = _annotation_node_ids(tree) - frozen_arguments = _frozen_argument_ids(tree) + exempt = _annotation_node_ids(tree) | _frozen_argument_ids(tree) | _typeddict_assigned_value_ids(tree) for node in ast.walk(tree): - if not isinstance(node, ast.expr) or id(node) in in_annotation or id(node) in frozen_arguments: + if not isinstance(node, ast.expr) or id(node) in exempt: continue kind = _construction_kind(node) if kind is None or node.lineno in comments.mutable_ok_lines: @@ -522,8 +623,10 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) f"mutable {kind}: this builds a collection that can be grown or rewritten. " f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator " f"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple " - f"/ ReadOnly TypedDict, or (if it really must be dynamic) a MappingProxyType " - f"wrapping a dict literal or comprehension (suppress: `# mutable-ok: `)", + f"/ ReadOnly TypedDict (a dict literal assigned to a name annotated with a " + f"same-module all-ReadOnly TypedDict is exempt), or (if it really must be dynamic) a " + f"MappingProxyType wrapping a dict literal or comprehension " + f"(suppress: `# mutable-ok: `)", ) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 2870a803db84..5bd6c8e5e654 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -193,6 +193,94 @@ def test_lit002_fix_message_names_mappingproxytype(tmp_path): assert "MappingProxyType" in messages[0] +_TYPEDDICT_PREFIX = ( + "from typing import Final, NotRequired, ReadOnly, TypedDict\n" + "class Td(TypedDict):\n" + " a: ReadOnly[int]\n" + " b: NotRequired[ReadOnly[str]]\n" +) + + +def test_dict_literal_annotated_as_typeddict_is_exempt(tmp_path): + assert "LIT002" not in _codes(tmp_path, _TYPEDDICT_PREFIX + "x: Td = {'a': 1}\n") + assert "LIT002" not in _codes(tmp_path, _TYPEDDICT_PREFIX + "y: Final[Td] = {'a': 1, 'b': 's'}\n") + + +def test_dict_literal_annotated_as_transitive_typeddict_subclass_is_exempt(tmp_path): + src = _TYPEDDICT_PREFIX + "class Sub(Td):\n c: ReadOnly[int]\nz: Final[Sub] = {'a': 1, 'c': 2}\n" + assert "LIT002" not in _codes(tmp_path, src) + + +def test_mutable_nested_inside_typeddict_literal_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "x: Td = {'a': 1, 'b': str([1])}\n") + + +def test_writable_typeddict_literal_gets_no_exemption(tmp_path): + src = "from typing import Final, TypedDict\nclass Wtd(TypedDict):\n a: int\nx: Final[Wtd] = {'a': 1}\n" + assert "LIT002" in _codes(tmp_path, src) + + +def test_typeddict_inheriting_writable_fields_gets_no_exemption(tmp_path): + loose_sub = _TYPEDDICT_PREFIX + "class Loose(Td):\n c: int\nz: Final[Loose] = {'a': 1, 'c': 2}\n" + assert "LIT002" in _codes(tmp_path, loose_sub) + writable_base = ( + "from typing import Final, ReadOnly, TypedDict\n" + "class W(TypedDict):\n a: int\n" + "class S(W):\n b: ReadOnly[int]\n" + "x: Final[S] = {'a': 1, 'b': 2}\n" + ) + assert "LIT002" in _codes(tmp_path, writable_base) + + +def test_typeddict_name_rebound_elsewhere_gets_no_exemption(tmp_path): + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "class Td:\n pass\nx: Td = {'a': 1}\n") + assert "LIT002" in _codes( + tmp_path, _TYPEDDICT_PREFIX + "def scope():\n Td = dict\n return Td\nx: Td = {'a': 1}\n" + ) + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "from elsewhere import Td\nx: Td = {'a': 1}\n") + assert "LIT002" in _codes( + tmp_path, _TYPEDDICT_PREFIX + "def scope(seq):\n for Td in seq:\n pass\nx: Td = {'a': 1}\n" + ) + assert "LIT002" in _codes( + tmp_path, _TYPEDDICT_PREFIX + "def scope(seq):\n Td, other = seq\n return other\nx: Td = {'a': 1}\n" + ) + assert "LIT002" in _codes( + tmp_path, _TYPEDDICT_PREFIX + "def scope(v):\n return (Td := v)\nx: Td = {'a': 1}\n" + ) + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "def scope(Td):\n return Td\nx: Td = {'a': 1}\n") + + +def test_nested_typeddict_gets_no_exemption(tmp_path): + assert "LIT002" in _codes( + tmp_path, + "from typing import ReadOnly, TypedDict\n" + "def scope():\n" + " class Td(TypedDict):\n" + " a: ReadOnly[int]\n" + "x: Td = {'a': 1}\n", + ) + + +def test_qualified_annotation_gets_no_exemption(tmp_path): + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "import elsewhere\nx: elsewhere.Td = {'a': 1}\n") + assert "LIT002" in _codes( + tmp_path, _TYPEDDICT_PREFIX + "import elsewhere\nx: Final[elsewhere.Td] = {'a': 1}\n" + ) + + +def test_star_import_disqualifies_typeddict_exemption(tmp_path): + assert "LIT002" in _codes(tmp_path, _TYPEDDICT_PREFIX + "from elsewhere import *\nx: Td = {'a': 1}\n") + + +def test_dict_literal_annotated_as_imported_or_unknown_type_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, "from other_module import Td\nx: Td = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\ny: Final = {'a': 1}\n") + + +def test_typeddict_call_spelling_is_not_construction(tmp_path): + assert "LIT002" not in _codes(tmp_path, _TYPEDDICT_PREFIX + "x: Final[Td] = Td(a=1)\n") + + def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): codes = _codes(tmp_path, "x: dict[str, int] = {} # mutable-ok: in-place buffer mutated hot path\n") assert "LIT001" not in codes