Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 126 additions & 2 deletions codegraph/analysis/lint.py
Original file line number Diff line number Diff line change
Expand Up @@ -318,10 +318,84 @@ def _check_sensitive_literal(
)


_TS_DB_PATTERN = r"\b(db|database)\.(select|insert|update|delete|execute|query)\b"
_PY_DB_PATTERN = r"\b(session|cursor|db)\.(query|execute|add|commit)\b"

# Higher-order iteration methods treated as implicit loop bodies.
_TS_ITER_METHODS = frozenset({"map", "forEach", "filter", "reduce", "flatMap"})
# Function-typed AST node types that can appear as callbacks.
_FUNC_NODE_TYPES = frozenset({"arrow_function", "function_expression", "function"})


def _py_chain_text(node: tree_sitter.Node, src: bytes) -> str | None:
"""Flatten a Python attribute-access chain to dotted text.

For a ``call`` node whose ``function`` child is an ``attribute``,
walks the attribute chain to produce e.g. ``session.execute``.
Returns ``None`` when the chain doesn't bottom out at an identifier.
"""
parts: list[str] = []
cur: tree_sitter.Node | None = node.child_by_field_name("function")
while cur is not None:
if cur.type == "attribute":
attr = cur.child_by_field_name("attribute")
if attr is not None:
parts.append(node_text(attr, src))
cur = cur.child_by_field_name("object")
elif cur.type in ("identifier",):
parts.append(node_text(cur, src))
return ".".join(reversed(parts))
else:
return None
return None


def _check_db_call_in_loop(
node: tree_sitter.Node, ctx: _LintContext
) -> LintFinding | None:
"""DB call inside a loop body — potential N+1 query."""
if ctx.is_test_file or ctx.loop_depth < 1:
return None

if ctx.language in _TS_LANGS:
if node.type != "call_expression" or not _is_chain_root(node):
return None
chain = _member_chain(node, ctx.src)
if chain is None:
return None
base, methods = chain
chain_text = ".".join([base, *methods])
pattern = str(ctx.rule.options.get("pattern") or _TS_DB_PATTERN)
elif ctx.language == "python":
if node.type != "call":
return None
chain_text = _py_chain_text(node, ctx.src) or ""
if not chain_text:
return None
pattern = str(ctx.rule.options.get("pattern") or _PY_DB_PATTERN)
else:
return None

if not re.search(pattern, chain_text):
return None

return LintFinding(
rule_id=ctx.rule.id,
severity=ctx.rule.severity,
message=_format_message(
ctx.rule, chain=chain_text, depth=ctx.loop_depth
),
file=ctx.rel_path,
line=node.start_point[0] + 1,
snippet=_snippet(node, ctx.src),
)


_CHECK_REGISTRY: dict[str, _Checker] = {
"console-in-prod": _check_console_in_prod,
"unfiltered-query": _check_unfiltered_query,
"sensitive-literal": _check_sensitive_literal,
"db-call-in-loop": _check_db_call_in_loop,
}


Expand All @@ -344,6 +418,12 @@ def _check_sensitive_literal(
severity="med",
message="Literal assigned to `{name}` looks sensitive ({reason})",
),
LintRule(
id="db-call-in-loop",
check="db-call-in-loop",
severity="high",
message="DB call `{chain}` inside a loop (N+1 risk)",
),
]


Expand Down Expand Up @@ -412,6 +492,45 @@ def _is_test_path(rel_path: str, language: str) -> bool:
return _is_ts_test_file(rel_path)


def _iter_children_with_depth(
node: tree_sitter.Node, base_depth: int, language: str, src: bytes
) -> list[tuple[tree_sitter.Node, int]]:
"""Return ``(child, depth)`` pairs for a node's children.

For TS/JS ``call_expression`` nodes whose member property is one of the
higher-order iteration methods (.map, .forEach, .filter, .reduce,
.flatMap), function-typed arguments receive ``base_depth + 1`` so the
DFS treats their bodies as implicit loop bodies. All other children
keep ``base_depth``.
"""
if language not in _TS_LANGS or node.type != "call_expression":
return [(child, base_depth) for child in reversed(node.children)]

fn = node.child_by_field_name("function")
if fn is None or fn.type != "member_expression":
return [(child, base_depth) for child in reversed(node.children)]

prop = fn.child_by_field_name("property")
if prop is None or node_text(prop, src) not in _TS_ITER_METHODS:
return [(child, base_depth) for child in reversed(node.children)]

args_node = node.child_by_field_name("arguments")
if args_node is None:
return [(child, base_depth) for child in reversed(node.children)]

result: list[tuple[tree_sitter.Node, int]] = []
for child in reversed(node.children):
if child == args_node:
for arg in reversed(child.children):
if arg.type in _FUNC_NODE_TYPES:
result.append((arg, base_depth + 1))
else:
result.append((arg, base_depth))
else:
result.append((child, base_depth))
return result


def _lint_file(
path: Path, rel_path: str, language: str, rules: list[LintRule]
) -> list[LintFinding]:
Expand Down Expand Up @@ -450,8 +569,13 @@ def _lint_file(
child_depth = (
loop_depth + 1 if node.type in _LOOP_NODE_TYPES else loop_depth
)
for child in reversed(node.children):
stack.append((child, child_depth))
# For .map/.forEach/.filter/.reduce/.flatMap calls (TS only), treat
# function-typed arguments as implicit loop bodies.
children_with_depth = _iter_children_with_depth(
node, child_depth, language, src
)
for child, cd in children_with_depth:
stack.append((child, cd))
return findings


Expand Down
18 changes: 15 additions & 3 deletions codegraph/parsers/python.py
Original file line number Diff line number Diff line change
Expand Up @@ -1319,6 +1319,11 @@ def _visit_nested_defs(
continue
stack.extend(node.children)

# Python loop node types that increase loop depth for their children.
_PY_LOOP_TYPES: frozenset[str] = frozenset(
{"for_statement", "while_statement"}
)

def _collect_calls(
self,
node: tree_sitter.Node,
Expand All @@ -1328,9 +1333,11 @@ def _collect_calls(
edges: list[Edge],
) -> None:
"""Walk subtree collecting call expressions, stopping at nested defs."""
stack: list[tree_sitter.Node] = list(node.children)
stack: list[tuple[tree_sitter.Node, int]] = [
(child, 0) for child in node.children
]
while stack:
child = stack.pop()
child, depth = stack.pop()
if child.type == "call":
func_child = child.child_by_field_name("function")
if func_child is None and child.children:
Expand All @@ -1352,6 +1359,8 @@ def _collect_calls(
"target_name": name,
"args": args,
"kwargs": kwargs,
"loop_depth": depth,
"in_loop": depth > 0,
},
))
# ``decorator`` subtrees are handled by ``_emit_decorator_calls``
Expand All @@ -1361,7 +1370,10 @@ def _collect_calls(
if child.type not in (
"class_definition", "function_definition", "decorator",
):
stack.extend(child.children)
child_depth = (
depth + 1 if child.type in self._PY_LOOP_TYPES else depth
)
stack.extend((gc, child_depth) for gc in child.children)

def _emit_decorator_calls(
self,
Expand Down
50 changes: 47 additions & 3 deletions codegraph/parsers/typescript.py
Original file line number Diff line number Diff line change
Expand Up @@ -1017,6 +1017,20 @@ def _handle_lexical_decl(
self._collect_calls(body, rel, func_id, src, edges)
self._collect_fetches(body, rel, func_id, src, nodes, edges)

# Higher-order iteration methods whose function-typed arguments are treated
# as implicit loop bodies for loop-depth tracking purposes.
_ITER_METHODS: frozenset[str] = frozenset(
{"map", "forEach", "filter", "reduce", "flatMap"}
)
# TS/JS loop node types that increase loop depth for their children.
_TS_LOOP_TYPES: frozenset[str] = frozenset(
{"for_statement", "for_in_statement", "while_statement", "do_statement"}
)
# Function-typed AST node types that can appear as callbacks.
_FUNC_ARG_TYPES: frozenset[str] = frozenset(
{"arrow_function", "function_expression", "function"}
)

def _collect_calls(
self,
node: tree_sitter.Node,
Expand All @@ -1025,9 +1039,11 @@ def _collect_calls(
src: bytes,
edges: list[Edge],
) -> None:
stack: list[tree_sitter.Node] = list(node.children)
stack: list[tuple[tree_sitter.Node, int]] = [
(child, 0) for child in node.children
]
while stack:
child = stack.pop()
child, depth = stack.pop()
if child.type == "call_expression":
func_child = child.child_by_field_name("function")
if func_child is None and child.children:
Expand Down Expand Up @@ -1056,9 +1072,37 @@ def _collect_calls(
"target_name": name,
"args": call_args,
"kwargs": call_kwargs,
"loop_depth": depth,
"in_loop": depth > 0,
},
))
stack.extend(child.children)
# Determine whether this call is an iteration method
# (.map/.forEach/.filter/.reduce/.flatMap) and bump depth
# for any function-typed arguments.
iter_method = ""
if func_child is not None and func_child.type == "member_expression":
prop_node = func_child.child_by_field_name("property")
if prop_node is not None:
iter_method = node_text(prop_node, src)
if iter_method in self._ITER_METHODS and args_node is not None:
# The callee subtree may itself contain calls
# (e.g. ``getUsers().map(cb)``) — keep visiting it
# at the current depth.
stack.extend(
(gc, depth) for gc in func_child.children
)
for arg in args_node.children:
if arg.type in self._FUNC_ARG_TYPES:
stack.extend(
(gc, depth + 1) for gc in arg.children
)
elif arg.type not in (",", "(", ")"):
stack.extend((gc, depth) for gc in arg.children)
continue
child_depth = (
depth + 1 if child.type in self._TS_LOOP_TYPES else depth
)
stack.extend((gc, child_depth) for gc in child.children)

# ------------------------------------------------------------------
# DF2: HTTP call-site detection (fetch / axios / SWR / api-clients)
Expand Down
19 changes: 19 additions & 0 deletions tests/fixtures/df0_typescript/loop_calls.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
// Fixture for loop_depth / in_loop metadata tests.

export function withLoops(items: string[]): void {
for (const item of items) {
doWork(item);
}

items.forEach((x) => {
doWork(x);
});
}

export function withoutLoop(): void {
doWork("direct");
}

export function chainedBeforeMap(): void {
getUsers().map((u) => doWork(u));
}
12 changes: 12 additions & 0 deletions tests/fixtures/lint_sample/src/batch.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
from db import session


# for loop with session.execute — should fire db-call-in-loop
def process_items(items):
for item in items:
session.execute("INSERT INTO log VALUES (?)", [item])


# non-loop call — must NOT fire
def setup():
session.execute("CREATE TABLE IF NOT EXISTS log (val TEXT)")
20 changes: 20 additions & 0 deletions tests/fixtures/lint_sample/src/loops.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import { db, jobs } from "./schema";

// for...of loop — should fire db-call-in-loop
export async function insertAll(items: string[]): Promise<void> {
for (const item of items) {
await db.insert(jobs).values(item);
}
}

// .forEach callback — should fire db-call-in-loop
export async function updateAll(ids: number[]): Promise<void> {
ids.forEach(async (id) => {
await db.update(jobs).set({ processed: true }).where(id);
});
}

// non-loop call — must NOT fire
export async function insertOne(item: string): Promise<void> {
await db.insert(jobs).values(item);
}
39 changes: 39 additions & 0 deletions tests/test_df0_python.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,3 +185,42 @@ def test_call_literal_kinds_captured(
_, edges = _run_extractor_on_source(tmp_path, src)
call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
assert call.metadata["args"] == [expected]


# --- loop_depth / in_loop metadata ----------------------------------------


def test_call_inside_for_loop_has_loop_depth(tmp_path: Path) -> None:
src = "def caller(items):\n for item in items:\n f(item)\n"
_, edges = _run_extractor_on_source(tmp_path, src)
call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
assert call.metadata["loop_depth"] == 1
assert call.metadata["in_loop"] is True


def test_call_inside_while_loop_has_loop_depth(tmp_path: Path) -> None:
src = "def caller(cond):\n while cond:\n f()\n"
_, edges = _run_extractor_on_source(tmp_path, src)
call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
assert call.metadata["loop_depth"] == 1
assert call.metadata["in_loop"] is True


def test_call_outside_loop_has_zero_loop_depth(tmp_path: Path) -> None:
src = "def caller():\n f()\n"
_, edges = _run_extractor_on_source(tmp_path, src)
call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
assert call.metadata["loop_depth"] == 0
assert call.metadata["in_loop"] is False


def test_nested_loop_increments_depth(tmp_path: Path) -> None:
src = (
"def caller(matrix):\n"
" for row in matrix:\n"
" for cell in row:\n"
" f(cell)\n"
)
_, edges = _run_extractor_on_source(tmp_path, src)
call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
assert call.metadata["loop_depth"] == 2
Loading
Loading