diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index d83a1a7512f..adbf44f15c8 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -206,18 +206,26 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, . # --------------------------------------------------------------------------- # -def mutable_names_in(annotation: ast.expr) -> Iterator[str]: - """Yield mutable-collection names anywhere inside an annotation expression. +class MutableRef(NamedTuple): + name: str + line: int + + +def mutable_names_in(annotation: ast.expr) -> Iterator[MutableRef]: + """Yield each mutable-collection name inside an annotation, with its source line. Matches bare names (`dict`, `MutableMapping`) and dotted access (`typing.Dict`, `collections.deque`, `collections.abc.MutableMapping`), descends through nesting (`Mapping[str, list[int]]`, `tuple[set[int], ...]`) and string forward references. + Each name carries its own line so a violation in a multi-line annotation is reported + where the name sits -- not on the annotation's first line -- and a `# mutable-ok` on + that line suppresses it. """ for node in ast.walk(annotation): if isinstance(node, ast.Name) and node.id in MUTABLE_COLLECTIONS: - yield node.id + yield MutableRef(node.id, node.lineno) elif isinstance(node, ast.Attribute) and node.attr in MUTABLE_COLLECTIONS: - yield node.attr + yield MutableRef(node.attr, node.lineno) elif isinstance(node, ast.Constant): value: object = node.value # forward references arrive as string constants if isinstance(value, str): @@ -225,7 +233,9 @@ def mutable_names_in(annotation: ast.expr) -> Iterator[str]: inner = ast.parse(value, mode="eval").body except SyntaxError: continue - yield from mutable_names_in(inner) + # The inner parse numbers lines from 1 inside the string, so anchor every + # name it yields to the forward-ref string's own line in the file. + yield from (MutableRef(ref.name, node.lineno) for ref in mutable_names_in(inner)) def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: @@ -240,11 +250,15 @@ def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: def _annotation_violations( - path: Path, annotation: ast.expr | None, line: int, where: str, ok_lines: frozenset[int] + path: Path, annotation: ast.expr | None, where: str, ok_lines: frozenset[int] ) -> Iterator[Violation]: - if annotation is None or line in ok_lines: + if annotation is None: return - yield from (_mutable_ann(path, line, name, where) for name in mutable_names_in(annotation)) + yield from ( + _mutable_ann(path, ref.line, ref.name, where) + for ref in mutable_names_in(annotation) + if ref.line not in ok_lines + ) def _function_violations( @@ -254,14 +268,14 @@ def _function_violations( args = node.args for arg in (*args.posonlyargs, *args.args, *args.kwonlyargs): yield from _annotation_violations( - path, arg.annotation, arg.lineno, f"parameter `{arg.arg}` of `{node.name}`", mutable_ok + path, arg.annotation, f"parameter `{arg.arg}` of `{node.name}`", mutable_ok ) # *args is allowed when typed (it's just a tuple); ruff ANN002 forces the # annotation, so here we only add the LIT001 mutable-collection check on the element type. if args.vararg is not None: yield from _annotation_violations( - path, args.vararg.annotation, args.vararg.lineno, f"`*args` of `{node.name}`", mutable_ok + path, args.vararg.annotation, f"`*args` of `{node.name}`", mutable_ok ) # **kwargs is banned outright (LIT008): it erases the keyword contract and forces @@ -278,7 +292,7 @@ def _function_violations( if node.returns is not None: yield from _annotation_violations( - path, node.returns, node.returns.lineno, f"return type of `{node.name}`", mutable_ok + path, node.returns, f"return type of `{node.name}`", mutable_ok ) @@ -293,8 +307,7 @@ def iter_annotation_violations(path: Path, tree: ast.AST, comments: Comments) -> elif isinstance(node, ast.AnnAssign): target = node.target.id if isinstance(node.target, ast.Name) else "" yield from _annotation_violations( - path, node.annotation, node.lineno, - f"the type of `{target}`", comments.mutable_ok_lines, + path, node.annotation, f"the type of `{target}`", comments.mutable_ok_lines, ) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 436904b017c..e49e196477f 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -25,6 +25,12 @@ def _codes(tmp_path, source): return [v.code for v in checker.check_file(f)] +def _lines(tmp_path, source, code): + f = tmp_path / "snippet.py" + f.write_text(source, encoding="utf-8") + return sorted(v.line for v in checker.check_file(f) if v.code == code) + + # --------------------------------------------------------------------------- # # Comment scanning (the readline path) — LIT003 / LIT004 / LIT005 # --------------------------------------------------------------------------- # @@ -132,6 +138,52 @@ def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): assert "LIT002" not in codes +def test_multiline_annotation_is_reported_at_the_name_not_the_first_line(tmp_path): + # The mutable name sits three lines below where the annotation opens. The violation must + # point at the name's own line so the message, the PR-diff gate's "introduced here" hint, + # and any `# mutable-ok` all land where the name actually is. + src = ( + "from collections.abc import Mapping\n" # 1 + "def f() -> Mapping[\n" # 2 + " str,\n" # 3 + " list[int],\n" # 4 <- the mutable name lives here + "]:\n" # 5 + " ...\n" # 6 + ) + assert _lines(tmp_path, src, "LIT001") == [4] + + +def test_mutable_ok_on_the_name_line_of_a_multiline_annotation_suppresses(tmp_path): + # Suppression must be honored on the line carrying the mutable name, not on the + # annotation's opening line; the latter is where a developer would never think to put it. + src = ( + "x: Mapping[\n" + " str,\n" + " list[int], # mutable-ok: in-place buffer mutated on the hot path\n" + "] = make()\n" + ) + assert "LIT001" not in _codes(tmp_path, src) + + +def test_mutable_ok_on_the_opening_line_no_longer_blankets_a_later_name(tmp_path): + # The opening line carries the suppression but the mutable name is two lines down, so the + # name is still flagged: suppression is per-name-line, never a blanket over the whole span. + src = ( + "x: dict[ # mutable-ok: only meant to cover this line\n" + " str,\n" + " list[int],\n" + "] = make()\n" + ) + assert _lines(tmp_path, src, "LIT001") == [3] + + +def test_forward_ref_violation_anchors_to_the_string_line(tmp_path): + # A forward-ref string is parsed on its own, numbering lines from 1 inside the quotes; + # the violation must still report the string's line in the file, not line 1. + src = "a = 1\nb = 2\nx: 'dict[str, int]'\n" + assert _lines(tmp_path, src, "LIT001") == [3] + + # --------------------------------------------------------------------------- # # Casts (LIT006) # --------------------------------------------------------------------------- #