Skip to content

Commit 1c7ec39

Browse files
juenglinjuenglin
committed
ci: fix two pool-props lint gaps found in CodeRabbit review
Finding 1 (minor): a type-annotated construction ``props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()`` produces an ast.AnnAssign, not ast.Assign, so the rule never saw the constructor. Extend the construction detection to cover AnnAssign. Finding 2 (major): when the same variable is assigned twice, ``props = CUmemPoolProps()`` # capped ``props = CUmemPoolProps()`` # uncapped, but lint was silent the old cap remained in the shared `capped` set and excused the second construction. Fix: process nodes in source order and track the latest uncapped construction per variable name. When the variable is reassigned, the old cap is reset; the new construction must have its own cap. Any construction still pending at end-of-scope is reported. Also change _walk_scope to yield nodes in source order (previously reversed) so the forward-pass cap-state tracking is order-correct. Adds targeted tests for both findings.
1 parent 069e168 commit 1c7ec39

2 files changed

Lines changed: 90 additions & 17 deletions

File tree

‎ci/tools/check_mempool_hygiene.py‎

Lines changed: 47 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -97,13 +97,14 @@ def _opted_out(lines: list[str], node: ast.AST) -> bool:
9797

9898

9999
def _walk_scope(scope: ast.AST) -> Iterator[ast.AST]:
100-
"""Yield ``scope``'s descendants, yielding but not entering nested function/lambda/class bodies."""
101-
pending = list(ast.iter_child_nodes(scope))
100+
"""Yield ``scope``'s descendants in source order, not entering nested function/lambda/class bodies."""
101+
# Push children reversed so the leftmost (earliest) child pops first.
102+
pending = list(reversed(list(ast.iter_child_nodes(scope))))
102103
while pending:
103104
node = pending.pop()
104105
yield node
105106
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.ClassDef)):
106-
pending.extend(ast.iter_child_nodes(node))
107+
pending.extend(reversed(list(ast.iter_child_nodes(node))))
107108

108109

109110
def _assign_targets(node: ast.AST) -> list[ast.expr]:
@@ -118,35 +119,64 @@ def _assign_targets(node: ast.AST) -> list[ast.expr]:
118119
def _pool_props_violations(path: Path, scope: ast.AST, lines: list[str]) -> list[Violation]:
119120
"""Report each ``var = PoolPropsType(...)`` with no ``var.max_size`` set in the same scope.
120121
121-
Recurses so a cap only counts in the scope that built the props; a cap on a
122-
same-named variable elsewhere does not excuse it.
122+
Processes nodes in source order. When ``var`` is reassigned to a new
123+
construction, any cap set on the previous object is reset; only a cap that
124+
follows the new construction (in the same scope) excuses it.
125+
126+
Recurses into nested scopes so a cap set in another function does not count.
127+
128+
Handles both plain and type-annotated assignments (``ast.Assign`` and
129+
``ast.AnnAssign``).
123130
"""
124131
found: list[Violation] = []
125-
constructed: list[tuple[ast.Assign, str, str]] = [] # (node, variable, type name)
126-
capped: set[str] = set()
132+
# Maps variable name -> (construction node, type name) for the latest
133+
# uncapped construction. Removed when the variable is capped; overwritten
134+
# (with a report of the old entry if it was still pending) on reassignment.
135+
pending: dict[str, tuple[ast.AST, str]] = {}
136+
127137
for node in _walk_scope(scope):
128138
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
129139
found += _pool_props_violations(path, node, lines)
130140
continue
141+
# Cap: props.max_size = ... or props.maxSize = ...
131142
for target in _assign_targets(node):
132143
if (
133144
isinstance(target, ast.Attribute)
134145
and target.attr in POOL_PROPS_CAP_ATTRIBUTES
135146
and isinstance(target.value, ast.Name)
136147
):
137-
capped.add(target.value.id)
138-
if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.value, ast.Call):
139-
(target,) = node.targets
140-
type_name = _callee_name(node.value)
141-
if isinstance(target, ast.Name) and type_name in POOL_PROPS_TYPES:
142-
constructed.append((node, target.id, type_name))
143-
for node, variable, type_name in constructed:
144-
if variable not in capped and not _opted_out(lines, node):
148+
pending.pop(target.value.id, None)
149+
# Construction: props = PoolPropsType(...) or props: T = PoolPropsType(...)
150+
con_target = con_value = None
151+
if isinstance(node, ast.Assign) and len(node.targets) == 1:
152+
con_target, con_value = node.targets[0], node.value
153+
elif isinstance(node, ast.AnnAssign) and node.value is not None:
154+
con_target, con_value = node.target, node.value
155+
if (
156+
isinstance(con_target, ast.Name)
157+
and isinstance(con_value, ast.Call)
158+
and _callee_name(con_value) in POOL_PROPS_TYPES
159+
):
160+
var, type_name = con_target.id, _callee_name(con_value)
161+
if var in pending:
162+
prev_node, prev_type = pending.pop(var)
163+
if not _opted_out(lines, prev_node):
164+
found.append(
165+
Violation(
166+
RULE_PROPS,
167+
f"{path.as_posix()}:{prev_node.lineno}: {prev_type} assigned to '{var}' "
168+
f"without {var}.max_size / {var}.maxSize",
169+
)
170+
)
171+
pending[var] = (node, type_name)
172+
173+
for var, (node, type_name) in pending.items():
174+
if not _opted_out(lines, node):
145175
found.append(
146176
Violation(
147177
RULE_PROPS,
148-
f"{path.as_posix()}:{node.lineno}: {type_name} assigned to '{variable}' "
149-
f"without {variable}.max_size / {variable}.maxSize",
178+
f"{path.as_posix()}:{node.lineno}: {type_name} assigned to '{var}' "
179+
f"without {var}.max_size / {var}.maxSize",
150180
)
151181
)
152182
return found

‎ci/tools/tests/test_check_mempool_hygiene.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,16 @@ def test_the_live_test_suite_is_clean():
103103
pytest.param("def f():\n props = cuda.CUmemPoolProps_v1()\n props.maxSize = POOL_SIZE\n", id="in-function"),
104104
# A field of an alloc-node params object is not a constructor call, so it is not checked.
105105
pytest.param("params.poolProps.allocType = TYPE\nparams.pool_props = other\n", id="graph-params-fields"),
106+
# Type-annotated assignment (ast.AnnAssign) must also be detected and treated as capped.
107+
pytest.param(
108+
"props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
109+
id="annotated-capped",
110+
),
111+
# Second construction caps correctly; the first was already capped.
112+
pytest.param(
113+
"props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
114+
id="two-constructions-both-capped",
115+
),
106116
]
107117

108118
POOL_PROPS_UNCAPPED = [
@@ -118,6 +128,13 @@ def test_the_live_test_suite_is_clean():
118128
"def make():\n return cuda.CUmemPoolProps()\n\n\ndef f():\n props = cuda.CUmemPoolProps()\n g(props)\n\n\ndef g(props):\n props.maxSize = POOL_SIZE\n",
119129
id="capped-in-another-function",
120130
),
131+
# Type-annotated assignment without a cap must also be reported (ast.AnnAssign, Finding 1).
132+
pytest.param("props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n", id="annotated-uncapped"),
133+
# A cap on the first object does not cover a second construction of the same name (Finding 2).
134+
pytest.param(
135+
"props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\n",
136+
id="reassignment-second-uncapped",
137+
),
121138
]
122139

123140

@@ -174,3 +191,29 @@ def test_main_reports_both_rules_separately(tmp_path, capsys):
174191
err = capsys.readouterr().err
175192
assert "memory pools created by tests must set max_size:" in err
176193
assert "pool properties created by tests must set max_size / maxSize" in err
194+
195+
196+
@pytest.mark.agent_authored(model="claude-sonnet-4.6")
197+
def test_annotated_pool_props_assignment_reports_correct_line(tmp_path):
198+
# Finding 1: props: T = CUmemPoolProps() is ast.AnnAssign; it must be detected.
199+
path = write(tmp_path, "x = 1\nprops: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n")
200+
201+
(violation,) = violations_in(path)
202+
203+
assert ":2:" in violation
204+
205+
206+
@pytest.mark.agent_authored(model="claude-sonnet-4.6")
207+
def test_reassignment_reports_only_the_uncapped_construction(tmp_path):
208+
# Finding 2: a cap on the first object must not excuse a second construction
209+
# of the same variable. Only the second (uncapped) object should be reported.
210+
source = (
211+
"props = cuda.CUmemPoolProps()\n" # line 1
212+
"props.maxSize = POOL_SIZE\n" # line 2 — caps line 1
213+
"props = cuda.CUmemPoolProps()\n" # line 3 — new object, no cap
214+
)
215+
path = write(tmp_path, source)
216+
217+
(violation,) = violations_in(path)
218+
219+
assert ":3:" in violation

0 commit comments

Comments
 (0)