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

* feat(lint): LIT013 caps comprehensions at one for and one if clause

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* docs(lint): rewrap the type discipline gate rule list

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): honor comprehension-ok on any line a comprehension spans

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): scope comprehension-ok to the innermost comprehension spanning it

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): break equal-span suppression ties toward the inner comprehension

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(lint): let single-line and only violating comprehensions own comprehension-ok markers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(lint): type tmp_path in LIT014 tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

AGENTS.md

@@ -96,6 +96,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
 - No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
   - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
   - Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
+  - Comprehensions take at most one `for` clause and one `if` clause (LIT014); split stacked clauses into a helper generator, a named intermediate, or a plain loop. Suppress with `# comprehension-ok: <reason>` only when unavoidable
 - Use dependency injection
 - Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
 - Use tagged unions + match

scripts/check_type_discipline.py

@@ -40,7 +40,8 @@ LIT003  noqa suppression without rule codes or without a reason.
 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.
@@ -103,6 +104,15 @@ LIT013  A `# <token>-ok: <reason>` suppression on a line where none of the rules
         that token suppresses fires. Like ruff's RUF100: a marker that suppresses
         nothing rots in place and hides real violations that land on the line
         later. Delete it.
+LIT014  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. The
+        marker belongs to the innermost violating comprehension spanning that
+        line, and also to any single-line violating comprehension on 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.
@@ -206,6 +216,7 @@ GUARD_OK_RE = re.compile(r"#\s*guard-ok(?::\s*(?P<reason>.*))?")
 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>.*))?")
 
 @dataclass(frozen=True, slots=True)
 class _OkToken:
@@ -224,6 +235,7 @@ OK_SUPPRESSIONS: Final[tuple[_OkToken, ...]] = (
     _OkToken("kwargs-ok", KWARGS_OK_RE, frozenset(("LIT008",))),
     _OkToken("rebind-ok", REBIND_OK_RE, frozenset(("LIT010", "LIT011"))),
     _OkToken("writable-ok", WRITABLE_OK_RE, frozenset(("LIT012",))),
+    _OkToken("comprehension-ok", COMPREHENSION_OK_RE, frozenset(("LIT014",))),
 )
 
 
@@ -1035,6 +1047,83 @@ def iter_typeddict_violations(path: Path, tree: ast.AST) -> Iterator[Violation]:
         )
 
 
+# --------------------------------------------------------------------------- #
+# Stacked comprehension clauses (LIT014)
+# --------------------------------------------------------------------------- #
+
+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 _clause_counts(node: ast.expr) -> tuple[int, int]:
+    return (
+        len(node.generators),
+        sum(len(g.ifs) for g in node.generators),
+    )
+
+
+def _violates(node: ast.expr) -> bool:
+    for_count, if_count = _clause_counts(node)
+    return for_count > 1 or if_count > 1
+
+
+def _comprehension_owners(tree: ast.AST, ok_lines: frozenset[int]) -> Mapping[int, int]:
+    """id(node) -> marker line for each `# comprehension-ok` line's owner.
+
+    Only violating comprehensions own markers. Each marker belongs to the
+    innermost violating comprehension whose span contains it (line span first,
+    column width breaks ties) plus every violating comprehension whose whole
+    span is that single line, so a comment inside a nested comprehension never
+    silences a multi-line enclosing one and a violation sharing its only line
+    can still be suppressed.
+    """
+    violating: Final = tuple(
+        n for n in ast.walk(tree) if isinstance(n, COMPREHENSION_NODES) and _violates(n)
+    )
+
+    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 owners(line: int) -> tuple[ast.expr, ...]:
+        containing: Final = tuple(n for n in violating if line in _span(n))
+        innermost: Final = min(containing, key=nesting_key, default=None)
+        single_line: Final = tuple(n for n in violating if len(_span(n)) == 1 and n.lineno == line)
+        return (*single_line, *(() if innermost is None else (innermost,)))
+
+    return MappingProxyType({id(o): line for line in ok_lines for o in owners(line)})
+
+
+def iter_comprehension_violations(
+    path: Path, tree: ast.AST, ok_lines: frozenset[int]
+) -> Iterator[tuple[Violation, bool]]:
+    """(violation, owned) pairs for every violating comprehension.
+
+    An owned comprehension reports at its marker's line so apply_suppressions
+    drops it and counts the marker as used; an unowned one reports at its own
+    line and is kept verbatim, since a marker suppresses only its owner even
+    when another violation shares that line.
+    """
+    owners: Final = _comprehension_owners(tree, ok_lines)
+    for node in ast.walk(tree):
+        if not isinstance(node, COMPREHENSION_NODES) or not _violates(node):
+            continue
+        for_count, if_count = _clause_counts(node)
+        yield (
+            Violation(
+                path,
+                owners.get(id(node), node.lineno),
+                "LIT014",
+                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>`)",
+            ),
+            id(node) in owners,
+        )
+
+
 # --------------------------------------------------------------------------- #
 # Suppression application and unused suppressions (LIT013)
 # --------------------------------------------------------------------------- #
@@ -1087,8 +1176,13 @@ def check_file(path: Path) -> tuple[Violation, ...]:
     except SyntaxError as exc:
         return (*violations, Violation(path, exc.lineno or 0, "LIT000", f"syntax error: {exc.msg}"))
 
+    comprehension_violations: Final = tuple(
+        iter_comprehension_violations(path, tree, suppressions["comprehension-ok"])
+    )
+
     return (
         *violations,
+        *(v for v, owned in comprehension_violations if not owned),
         *apply_suppressions(
             path,
             (
@@ -1099,6 +1193,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
                 *iter_final_violations(path, tree),
                 *iter_param_violations(path, tree),
                 *iter_typeddict_violations(path, tree),
+                *(v for v, owned in comprehension_violations if owned),
             ),
             suppressions,
         ),

scripts/type_discipline_gate.py

@@ -15,8 +15,12 @@ without codes or reason), LIT006 (cast), LIT008 (`**kwargs`), LIT009 (inert
 (assignment without a Final declaration; suppress deliberate rebinding with
 `# rebind-ok: <reason>`), LIT011 (parameter rebinding or in-place mutation), and
 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 LIT014 (comprehension with more than one `for`
+or `if` clause; suppress with `# comprehension-ok: <reason>` on a spanned
+line, which belongs to the innermost violating comprehension spanning it and
+to any single-line violating comprehension on that line) 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; LIT013 (`*-ok` suppression
 that suppresses nothing) is frozen at 0 for the same reason; and LIT007
 (TypeGuard/TypeIs) is a hard zero.
@@ -198,7 +202,8 @@ def cmd_check(base: str) -> None:
         "Remove the new violations, give each a reason (`# noqa: XXX  # <reason>`, "
         "`# pyright: ignore[rule]  # <reason>`, `# mutable-ok: <reason>`, "
         "`# cast-ok: <reason>`, `# guard-ok: <reason>`, `# kwargs-ok: <reason>`, "
-        "`# rebind-ok: <reason>`, `# writable-ok: <reason>`), or remove an equal "
+        "`# rebind-ok: <reason>`, `# writable-ok: <reason>`, "
+        "`# comprehension-ok: <reason>`), or remove an equal "
         "number elsewhere; the ceiling "
         "is the limit in type-discipline-budget.json."
     )

tests/test_litellm/test_check_type_discipline.py

@@ -685,6 +685,169 @@ def test_writable_ok_without_reason_is_lit005_and_does_not_suppress(tmp_path):
     assert "LIT012" in codes
 
 
+# --------------------------------------------------------------------------- #
+# Stacked comprehension clauses (LIT014)
+# --------------------------------------------------------------------------- #
+
+
+def test_two_for_clauses_are_flagged(tmp_path: Path):
+    assert "LIT014" 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: Path):
+    assert "LIT014" 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: Path):
+    assert "LIT014" 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: Path):
+    assert "LIT014" 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: Path):
+    assert "LIT014" in _codes(tmp_path, "d = {k: v for a in xs for k, v in a}\n")
+    assert "LIT014" in _codes(tmp_path, "s = {x for a in xs for x in a}\n")
+    assert "LIT014" 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: Path):
+    assert "LIT014" not in _codes(tmp_path, "y = [[v for v in a] for a in xs]\n")
+
+
+def test_comprehension_ok_with_reason_suppresses_lit014(tmp_path: 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 "LIT014" not in codes
+
+
+def test_comprehension_ok_on_any_spanned_line_suppresses_lit014(tmp_path: 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 "LIT014" not in _codes(tmp_path, src)
+
+
+def test_comprehension_ok_after_the_closing_line_does_not_suppress(tmp_path: 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"
+    )
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    violations = checker.check_file(f)
+    assert [v.line for v in violations if v.code == "LIT014"] == [1]
+    assert [v.line for v in violations if v.code == "LIT013"] == [5]
+
+
+def test_comprehension_ok_on_a_compliant_comprehension_is_an_unused_marker(tmp_path: Path):
+    f = tmp_path / "snippet.py"
+    f.write_text(
+        "y = tuple(x for x in xs if x)  # comprehension-ok: kept for readability\n",
+        encoding="utf-8",
+    )
+    violations = checker.check_file(f)
+    assert [v.line for v in violations if v.code == "LIT013"] == [1]
+    assert "LIT014" not in [v.code for v in violations]
+
+
+def test_comprehension_ok_without_reason_is_lit005_and_does_not_suppress(tmp_path: Path):
+    codes = _codes(tmp_path, "y = [x for a in xs for x in a]  # comprehension-ok\n")
+    assert "LIT005" in codes
+    assert "LIT014" in codes
+
+
+def test_suppression_inside_inner_comprehension_does_not_silence_the_outer(tmp_path: 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")
+    violations = checker.check_file(f)
+    assert [v.line for v in violations if v.code == "LIT014"] == [1]
+    assert [v.code for v in violations if v.code == "LIT013"] == []
+
+
+def test_suppression_on_outer_closing_line_does_not_silence_the_inner(tmp_path: 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 == "LIT014"]
+    assert [v.line for v in flagged] == [3]
+
+
+def test_equal_span_marker_suppresses_every_violating_comprehension_on_its_line(tmp_path: 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")
+    violations = checker.check_file(f)
+    assert "LIT014" not in [v.code for v in violations]
+    assert "LIT013" not in [v.code for v in violations]
+
+
+def test_single_line_outer_with_violating_inner_is_suppressed(tmp_path: Path):
+    src = "y = [x for a in [z for i in ys for z in i] for x in a]  # comprehension-ok: nested flatten is fine\n"
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    violations = checker.check_file(f)
+    assert "LIT014" not in [v.code for v in violations]
+    assert "LIT013" not in [v.code for v in violations]
+
+
+def test_one_marker_suppresses_two_violating_sibling_comprehensions_on_its_line(tmp_path: Path):
+    src = "y = [x for a in xs for x in a] + [x for a in ys for x in a]  # comprehension-ok: paired flattens\n"
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    violations = checker.check_file(f)
+    assert "LIT014" not in [v.code for v in violations]
+    assert "LIT013" not in [v.code for v in violations]
+
+
+def test_marker_on_a_non_violating_inner_line_suppresses_the_violating_outer(tmp_path: Path):
+    src = (
+        "y = [\n"
+        "    x\n"
+        "    for a in [z for z in ys if z]  # comprehension-ok: flatten stays readable\n"
+        "    for x in a\n"
+        "]\n"
+    )
+    f = tmp_path / "snippet.py"
+    f.write_text(src, encoding="utf-8")
+    violations = checker.check_file(f)
+    assert "LIT014" not in [v.code for v in violations]
+    assert "LIT013" not in [v.code for v in violations]
+
+
+def test_violation_message_names_the_clause_counts(tmp_path: 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 == "LIT014"]
+    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

@@ -37,5 +37,8 @@
   },
   "LIT013": {
     "limit": 0
+  },
+  "LIT014": {
+    "limit": 369
   }
 }