diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index ce9eb391d55..0706c8a7bd8 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -18,13 +18,23 @@ LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehens Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`). Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a generator (`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / - NamedTuple / ReadOnly TypedDict, or (if it really must be dynamic) a - MappingProxyType wrapping a dict literal or comprehension. Generator expressions - and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`, + NamedTuple, a TypedDict-annotated dict literal, or (if it really must be + dynamic) a MappingProxyType wrapping a dict literal or comprehension. Generator + expressions and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`, `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. A dict literal whose assignment is + annotated with a TypedDict (`x: Final[MyTD] = {...}`; bare `x: Final = {...}` + does not qualify) is a fixed-shape build basedpyright checks key-by-key against + fields LIT012 keeps ReadOnly, not a growable accumulator, so it is exempt along + with the dict literals nested in it (nested TypedDict fields); any other + construction inside still counts. Detection is name-based: Final/ClassVar/ + Optional (and Annotated's first argument) unwrap, a PEP 604 union + (`MyTD | None`) qualifies through either arm, and any remaining named head + outside the mutable collections and Mapping/Any/object is taken to be a + TypedDict, since a dict literal assigned to any other named type would not + survive basedpyright. 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. @@ -138,6 +148,14 @@ MUTABLE_CONSTRUCTORS = frozenset(( # qualified `collections.deque(...)` still counts. QUALIFIED_CONSTRUCTORS = MUTABLE_CONSTRUCTORS - frozenset(("dict", "list", "set")) FREEZING_WRAPPERS = frozenset(("tuple", "frozenset", "MappingProxyType")) +# Wrappers unwrapped when deciding whether an assignment's annotation names a +# TypedDict (the LIT002 dict-literal exemption); bare, they name no type. Annotated +# is handled separately: only its first argument is type syntax. +TYPEDDICT_ANNOTATION_WRAPPERS = frozenset(("Final", "ClassVar", "Optional")) +# Heads that can type a dict literal without being a TypedDict. Every other named +# head counts as one: a dict literal assigned to any other named type would not +# survive basedpyright, which is the second gate behind this name-based check. +NON_TYPEDDICT_HEADS = MUTABLE_COLLECTIONS | frozenset(("Mapping", "Any", "object")) UNSAFE_GUARDS = frozenset(("TypeGuard", "TypeIs")) READONLY_QUALIFIER = "ReadOnly" # Qualifiers ReadOnly may nest under, in any order (PEP 705); for Annotated only the @@ -270,6 +288,14 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, . # --------------------------------------------------------------------------- # +def _head_name(node: ast.expr) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + def _is_literal_subscript(node: ast.AST) -> bool: if not isinstance(node, ast.Subscript): return False @@ -485,6 +511,61 @@ def _frozen_argument_ids(tree: ast.AST) -> frozenset[int]: ) +def _is_typeddict_annotation(annotation: ast.expr) -> bool: + """True iff the annotation names a TypedDict, by the name-based heuristic. + + Final/ClassVar/Optional unwrap (as does Annotated's first argument, the only + one that is type syntax), a PEP 604 union qualifies through either arm, string + forward references are parsed, and whatever named head remains counts as a + TypedDict unless it is a mutable collection or Mapping/Any/object -- the heads + that can type a dict literal without being one. Bare wrappers + (`x: Final = ...`) name no type and never qualify. + """ + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + try: + inner = ast.parse(annotation.value, mode="eval").body + except SyntaxError: + return False + return _is_typeddict_annotation(inner) + if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr): + return _is_typeddict_annotation(annotation.left) or _is_typeddict_annotation(annotation.right) + if isinstance(annotation, ast.Subscript): + head = _head_name(annotation.value) + if head in TYPEDDICT_ANNOTATION_WRAPPERS: + return _is_typeddict_annotation(annotation.slice) + if head == "Annotated": + first = annotation.slice.elts[0] if isinstance(annotation.slice, ast.Tuple) and annotation.slice.elts else None + return first is not None and _is_typeddict_annotation(first) + return head is not None and head not in NON_TYPEDDICT_HEADS + name = _head_name(annotation) + return ( + name is not None + and name not in NON_TYPEDDICT_HEADS + and name not in TYPEDDICT_ANNOTATION_WRAPPERS + and name != "Annotated" + ) + + +def _typeddict_build_ids(tree: ast.AST) -> frozenset[int]: + """ids() of every dict literal built under a TypedDict-annotated assignment. + + `x: Final[MyTD] = {...}` is a fixed-shape build: basedpyright checks each key + against the declared fields, which LIT012 keeps ReadOnly, so nothing here is + the seed-then-mutate accumulator LIT002 hunts. Dict literals nested in the + value (nested TypedDict fields) share the exemption; any other construction + inside it still counts, and a bare `x: Final = {...}` stays flagged. + """ + return frozenset( + id(sub) + for node in ast.walk(tree) + if isinstance(node, ast.AnnAssign) + and isinstance(node.value, ast.Dict) + and _is_typeddict_annotation(node.annotation) + for sub in ast.walk(node.value) + if isinstance(sub, ast.Dict) + ) + + def _construction_kind(node: ast.expr) -> str | None: """Human label if `node` builds a mutable collection, else None.""" if isinstance(node, ast.List): @@ -511,8 +592,14 @@ 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) + typeddict_builds = _typeddict_build_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 in_annotation + or id(node) in frozen_arguments + or id(node) in typeddict_builds + ): continue kind = _construction_kind(node) if kind is None or node.lineno in comments.mutable_ok_lines: @@ -521,9 +608,10 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) path, node.lineno, "LIT002", 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"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple, " + f"a TypedDict-annotated dict literal (`x: Final[MyTD] = {{...}}`), or (if it " + f"really must be dynamic) a MappingProxyType wrapping a dict literal or " + f"comprehension (suppress: `# mutable-ok: `)", ) @@ -851,14 +939,6 @@ def iter_param_violations(path: Path, tree: ast.AST, comments: Comments) -> Iter # --------------------------------------------------------------------------- # -def _head_name(node: ast.expr) -> str | None: - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - return node.attr - return None - - def _base_names(cls: ast.ClassDef) -> frozenset[str]: """The names of a class's bases; a subscripted base (`Foo[int]`) counts as `Foo`.""" return frozenset( diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 2870a803db8..84dd547ad80 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -199,6 +199,55 @@ def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): assert "LIT002" not in codes +def test_typeddict_annotated_dict_literal_is_exempt(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final\nfrom foo import MyTD\nx: Final[MyTD] = {'a': 1}\n" + ) + assert "LIT002" not in _codes(tmp_path, "from foo import MyTD\nx: MyTD = {'a': 1}\n") + assert "LIT002" not in _codes(tmp_path, "from typing import Final\nx: Final['MyTD'] = {'a': 1}\n") + assert "LIT002" not in _codes(tmp_path, "import foo\nfrom typing import Final\nx: Final[foo.MyTD] = {'a': 1}\n") + + +def test_wrapped_typeddict_annotations_share_the_exemption(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final, Optional\nx: Final[Optional[MyTD]] = {'a': 1}\n" + ) + assert "LIT002" not in _codes( + tmp_path, "from typing import Annotated, Final\nx: Final[Annotated[MyTD, 'meta']] = {'a': 1}\n" + ) + assert "LIT002" not in _codes( + tmp_path, "from typing import ClassVar\nclass C:\n x: ClassVar[MyTD] = {'a': 1}\n" + ) + assert "LIT002" not in _codes(tmp_path, "from typing import Final\nx: Final[MyTD | None] = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[dict[str, int] | None] = {'a': 1}\n") + + +def test_bare_final_dict_literal_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import ClassVar\nclass C:\n x: ClassVar = {'a': 1}\n") + + +def test_non_typeddict_annotations_do_not_exempt(tmp_path): + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[dict[str, int]] = {'a': 1}\n") + assert "LIT002" in _codes( + tmp_path, "from collections.abc import Mapping\nfrom typing import Final\nx: Final[Mapping[str, int]] = {'a': 1}\n" + ) + assert "LIT002" in _codes(tmp_path, "from typing import Any, Final\nx: Final[Any] = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[object] = {'a': 1}\n") + + +def test_typeddict_exemption_covers_only_dict_literals(tmp_path): + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = dict(a=1)\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = {k: 1 for k in ('a',)}\n") + + +def test_nested_dict_literals_share_the_typeddict_exemption(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final\nx: Final[Outer] = {'inner': {'a': 1}, 'steps': ({'b': 2},)}\n" + ) + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[Outer] = {'tags': ['a']}\n") + + # --------------------------------------------------------------------------- # # Casts (LIT006) # --------------------------------------------------------------------------- # diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8fbf9b291d5..8e55b1533ea 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22938 }, "LIT002": { - "limit": 27139 + "limit": 26901 }, "LIT003": { "limit": 269