feat(lint): exempt TypedDict-annotated dict literals from LIT002

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: <reason>`.
+        (`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: <reason>`.
 LIT003  noqa suppression without rule codes or without a reason.
         Required shape: `# noqa: TID251  # <reason>`
 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: <reason>`)",
+            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: <reason>`)",
         )
 
 

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