feat(lint): cap comprehensions at one for and one if clause (LIT014)

scripts/check_type_discipline.py

@@ -40,7 +40,8 @@
 LIT004  pyright/mypy ignore without bracketed codes or without a reason.
         Required shape: `# pyright: ignore[reportArgumentType]  # <reason>`
 LIT005  A `# mutable-ok` / `# cast-ok` / `# guard-ok` / `# kwargs-ok` /
-        `# rebind-ok` / `# writable-ok` suppression without a reason.
+        `# rebind-ok` / `# writable-ok` / `# comprehension-ok` suppression
+        without a reason.
 LIT006  `cast(...)` call. typing.cast is an unchecked assertion (the moral equivalent
         of TypeScript's `as`); it lies to the type checker with zero runtime guarantee.
         Validate into a concrete frozen type at the boundary instead.
@@ -99,6 +100,15 @@
         the functional form (`X = TypedDict("X", {...})`) is checked too. A base
         imported from another module is out of reach without import resolution.
         Suppress with `# writable-ok: <reason>`.
+LIT013  Comprehension with more than one `for` clause or more than one `if` clause,
+        in any of the four forms (list, set, dict, generator expression). Stacked
+        `for`s and `if`s read as nested loops and guards squashed onto one line;
+        split the comprehension into a helper generator, a named intermediate, or
+        a plain loop instead. A comprehension nested inside another's element or
+        iterable is its own node and is judged separately. Suppress with
+        `# comprehension-ok: <reason>` on any line the comprehension spans; when
+        comprehensions nest, the comment belongs to the innermost one spanning
+        that line.
 
 LIT000  Setup failure: a target file could not be read, or contains a syntax error.
         Reported as a violation rather than crashing the run.
@@ -182,6 +192,7 @@
 KWARGS_OK_RE = re.compile(r"#\s*kwargs-ok(?::\s*(?P<reason>.*))?")
 REBIND_OK_RE = re.compile(r"#\s*rebind-ok(?::\s*(?P<reason>.*))?")
 WRITABLE_OK_RE = re.compile(r"#\s*writable-ok(?::\s*(?P<reason>.*))?")
+COMPREHENSION_OK_RE = re.compile(r"#\s*comprehension-ok(?::\s*(?P<reason>.*))?")
 
 # Suppression tokens that must each carry a reason (LIT005).
 OK_SUPPRESSIONS: tuple[tuple[str, re.Pattern[str]], ...] = (
@@ -191,6 +202,7 @@
     ("kwargs-ok", KWARGS_OK_RE),
     ("rebind-ok", REBIND_OK_RE),
     ("writable-ok", WRITABLE_OK_RE),
+    ("comprehension-ok", COMPREHENSION_OK_RE),
 )
 
 
@@ -214,6 +226,7 @@ class Comments:
     kwargs_ok_lines: frozenset[int]
     rebind_ok_lines: frozenset[int]
     writable_ok_lines: frozenset[int]
+    comprehension_ok_lines: frozenset[int]
 
 
 # --------------------------------------------------------------------------- #
@@ -269,7 +282,7 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, .
         # tokenize raises TokenError (EOF mid-construct) or a SyntaxError subclass
         # (IndentationError / TabError) on malformed source; defer to ast.parse below,
         # which re-raises and is reported as LIT000 rather than crashing the run.
-        return Comments(frozenset(), frozenset(), frozenset(), frozenset(), frozenset(), frozenset()), ()
+        return Comments(frozenset(), frozenset(), frozenset(), frozenset(), frozenset(), frozenset(), frozenset()), ()
 
     def _lines_with(regex: re.Pattern[str]) -> frozenset[int]:
         return frozenset(line for line, text in comment_toks if _valid_ok(regex, text))
@@ -282,6 +295,7 @@ def _lines_with(regex: re.Pattern[str]) -> frozenset[int]:
             kwargs_ok_lines=_lines_with(KWARGS_OK_RE),
             rebind_ok_lines=_lines_with(REBIND_OK_RE),
             writable_ok_lines=_lines_with(WRITABLE_OK_RE),
+            comprehension_ok_lines=_lines_with(COMPREHENSION_OK_RE),
         ),
         tuple(v for line, text in comment_toks for v in _comment_violations(path, line, text)),
     )
@@ -1033,6 +1047,57 @@ def iter_typeddict_violations(path: Path, tree: ast.AST, comments: Comments) ->
         )
 
 
+# --------------------------------------------------------------------------- #
+# Stacked comprehension clauses (LIT013)
+# --------------------------------------------------------------------------- #
+
+COMPREHENSION_NODES = (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)
+
+
+def _span(node: ast.expr) -> range:
+    return range(node.lineno, (node.end_lineno or node.lineno) + 1)
+
+
+def _suppressed_comprehensions(tree: ast.AST, ok_lines: frozenset[int]) -> frozenset[int]:
+    """ids() of the comprehension each `# comprehension-ok` line suppresses.
+
+    A suppression belongs to the innermost comprehension whose span contains its
+    line, so a comment inside a nested comprehension never silences the enclosing
+    one (LIT013 judges each comprehension as its own node).
+    """
+    comps: Final = tuple(n for n in ast.walk(tree) if isinstance(n, COMPREHENSION_NODES))
+
+    def nesting_key(node: ast.expr) -> tuple[int, int]:
+        return (len(_span(node)), (node.end_col_offset or node.col_offset) - node.col_offset)
+
+    def owner(line: int) -> ast.expr | None:
+        containing: Final = tuple(n for n in comps if line in _span(n))
+        return min(containing, key=nesting_key, default=None)
+
+    return frozenset(id(o) for o in (owner(line) for line in ok_lines) if o is not None)
+
+
+def iter_comprehension_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]:
+    suppressed: Final = _suppressed_comprehensions(tree, comments.comprehension_ok_lines)
+    for node in ast.walk(tree):
+        if not isinstance(node, COMPREHENSION_NODES):
+            continue
+        for_count = len(node.generators)
+        if_count = sum(len(g.ifs) for g in node.generators)
+        if for_count <= 1 and if_count <= 1:
+            continue
+        if id(node) in suppressed:
+            continue
+        yield Violation(
+            path,
+            node.lineno,
+            "LIT013",
+            f"comprehension with {for_count} `for` clauses and {if_count} `if` clauses: "
+            f"at most one of each is allowed. Split it into a helper generator, a named "
+            f"intermediate, or a plain loop (suppress: `# comprehension-ok: <reason>`)",
+        )
+
+
 # --------------------------------------------------------------------------- #
 # Driver
 # --------------------------------------------------------------------------- #
@@ -1060,6 +1125,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
         *iter_final_violations(path, tree, comments),
         *iter_param_violations(path, tree, comments),
         *iter_typeddict_violations(path, tree, comments),
+        *iter_comprehension_violations(path, tree, comments),
     )
 
 

scripts/type_discipline_gate.py

@@ -13,10 +13,12 @@
 without codes or reason), LIT006 (cast), LIT008 (`**kwargs`), LIT009 (inert
 `# type: ignore`, dead syntax while enableTypeIgnoreComments is false), LIT010
 (assignment without a Final declaration; suppress deliberate rebinding with
-`# rebind-ok: <reason>`), LIT011 (parameter rebinding or in-place mutation), and
+`# rebind-ok: <reason>`), LIT011 (parameter rebinding or in-place mutation),
 LIT012 (TypedDict field without a `ReadOnly[...]` qualifier; suppress with
-`# writable-ok: <reason>`) carry limits at or above their current count to
-ratchet down; LIT005 (`*-ok` suppression without a reason) is frozen at limit 0
+`# writable-ok: <reason>`), and LIT013 (comprehension with more than one `for`
+or `if` clause; suppress with `# comprehension-ok: <reason>`) carry limits at
+or above their current count to ratchet down; LIT005 (`*-ok` suppression
+without a reason) is frozen at limit 0
 so any net-new reasonless suppression trips the gate; and LIT007
 (TypeGuard/TypeIs) is a hard zero.
 LIT010 and LIT011 were seeded at 1.5x the count left after the sweep that

tests/test_litellm/test_check_type_discipline.py

@@ -687,6 +687,120 @@ def test_writable_ok_without_reason_is_lit005_and_does_not_suppress(tmp_path):
     assert "LIT012" in codes
 
 
+# --------------------------------------------------------------------------- #
+# Stacked comprehension clauses (LIT013)
+# --------------------------------------------------------------------------- #
+
+
+def test_two_for_clauses_are_flagged(tmp_path):
+    assert "LIT013" in _codes(tmp_path, "y = [x for a in xs for x in a]\n")
+
+
+def test_two_ifs_on_one_generator_are_flagged(tmp_path):
+    assert "LIT013" in _codes(tmp_path, "y = [x for x in xs if x if x > 1]\n")
+
+
+def test_one_if_on_each_of_two_generators_is_flagged(tmp_path):
+    assert "LIT013" in _codes(tmp_path, "y = [x for a in xs if a for x in a if x]\n")
+
+
+def test_one_for_and_one_if_is_clean(tmp_path):
+    assert "LIT013" not in _codes(tmp_path, "y = tuple(x for x in xs if x)\n")
+
+
+def test_dict_set_and_generator_two_fors_are_each_flagged(tmp_path):
+    assert "LIT013" in _codes(tmp_path, "d = {k: v for a in xs for k, v in a}\n")
+    assert "LIT013" in _codes(tmp_path, "s = {x for a in xs for x in a}\n")
+    assert "LIT013" in _codes(tmp_path, "g = (x for a in xs for x in a)\n")
+
+
+def test_nested_comprehension_in_element_is_judged_separately(tmp_path):
+    assert "LIT013" not in _codes(tmp_path, "y = [[v for v in a] for a in xs]\n")
+
+
+def test_comprehension_ok_with_reason_suppresses_lit013(tmp_path):
+    codes = _codes(
+        tmp_path,
+        "y = [x for a in xs for x in a]  # comprehension-ok: flattens a stream of pairs, hot path\n",
+    )
+    assert "LIT013" not in codes
+
+
+def test_comprehension_ok_on_any_spanned_line_suppresses_lit013(tmp_path):
+    src = (
+        "y = [\n"
+        "    x for a in xs\n"
+        "    for x in a\n"
+        "]  # comprehension-ok: cartesian product is the clearest form\n"
+    )
+    assert "LIT013" not in _codes(tmp_path, src)
+
+
+def test_comprehension_ok_after_the_closing_line_does_not_suppress(tmp_path):
+    src = (
+        "y = [\n"
+        "    x for a in xs\n"
+        "    for x in a\n"
+        "]\n"
+        "# comprehension-ok: cartesian product is the clearest form\n"
+    )
+    assert "LIT013" in _codes(tmp_path, src)
+
+
+def test_comprehension_ok_without_reason_is_lit005_and_does_not_suppress(tmp_path):
+    codes = _codes(tmp_path, "y = [x for a in xs for x in a]  # comprehension-ok\n")
+    assert "LIT005" in codes
+    assert "LIT013" in codes
+
+
+def test_suppression_inside_inner_comprehension_does_not_silence_the_outer(tmp_path):
+    src = (
+        "y = [\n"
+        "    x\n"
+        "    for a in [\n"
+        "        z for i in ys\n"
+        "        for z in i\n"
+        "    ]  # comprehension-ok: inner flatten is the clearest form\n"
+        "    for x in a\n"
+        "]\n"
+    )
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    flagged = [v for v in checker.check_file(f) if v.code == "LIT013"]
+    assert [v.line for v in flagged] == [1]
+
+
+def test_suppression_on_outer_closing_line_does_not_silence_the_inner(tmp_path):
+    src = (
+        "y = [\n"
+        "    x\n"
+        "    for a in [z for i in ys for z in i]\n"
+        "    for x in a\n"
+        "]  # comprehension-ok: outer flatten is the clearest form\n"
+    )
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    flagged = [v for v in checker.check_file(f) if v.code == "LIT013"]
+    assert [v.line for v in flagged] == [3]
+
+
+def test_equal_span_suppression_belongs_to_the_inner_comprehension(tmp_path):
+    src = "y = [x for a in [z for i in ys for z in i] if a if x]  # comprehension-ok: inner flatten is fine\n"
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    flagged = [v for v in checker.check_file(f) if v.code == "LIT013"]
+    assert len(flagged) == 1
+    assert "1 `for` clauses and 2 `if` clauses" in flagged[0].message
+
+
+def test_violation_message_names_the_clause_counts(tmp_path):
+    f = tmp_path / "snippet.py"
+    f.write_text("y = [x for a in xs for x in a if x]\n", encoding="utf-8")
+    messages = [v.message for v in checker.check_file(f) if v.code == "LIT013"]
+    assert len(messages) == 1
+    assert "2 `for` clauses and 1 `if` clause" in messages[0]
+
+
 # --------------------------------------------------------------------------- #
 # Budget integrity: every emittable LIT rule (bar the LIT000 read/parse error) is gated
 # --------------------------------------------------------------------------- #

type-discipline-budget.json

@@ -34,5 +34,8 @@
   },
   "LIT012": {
     "limit": 4486
+  },
+  "LIT013": {
+    "limit": 363
   }
 }