From 8f9129618a4104fde401c3c080ced51af232217d Mon Sep 17 00:00:00 2001 From: Hugo Date: Fri, 4 Sep 2026 06:18:42 +0200 Subject: [PATCH 1/2] test: define conditional expressions, comprehensions, lambdas, sets, globals and chained assignments --- tests/test_cache.py | 2 +- tests/test_discovery.py | 22 +++ tests/test_django_http_clients.py | 2 +- tests/test_interprocedural.py | 4 +- tests/test_parallel.py | 2 +- tests/test_project.py | 2 +- tests/test_syntax_control_flow.py | 269 ++++++++++++++++++++++++++++++ tests/test_syntax_coverage.py | 10 +- 8 files changed, 302 insertions(+), 11 deletions(-) create mode 100644 tests/test_syntax_control_flow.py diff --git a/tests/test_cache.py b/tests/test_cache.py index d2e172c..21744af 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -183,7 +183,7 @@ def test_cached_modules_still_serve_project_plugins_and_correlation(tmp_path: Pa def test_unsupported_and_syntax_notes_are_cached_too(tmp_path: Path) -> None: - root = project(tmp_path / "src", {"nested.py": "def outer():\n def inner():\n pass\n return inner\n"}) + root = project(tmp_path / "src", {"nested.py": "def outer():\n class Inner:\n pass\n return Inner\n"}) cache = ProjectCache(tmp_path / "cache") first = engine.analyze_project(root, [PLUGINS], cache=cache) diff --git a/tests/test_discovery.py b/tests/test_discovery.py index ca4e65e..28e8699 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -94,3 +94,25 @@ def test_utf16_dependency_files_are_read(tmp_path: Path) -> None: assert analysis.dependencies.names == ("flask", "pyyaml") assert sorted(f.rule_id for f in analysis.findings) == ["vulnerable-dependency", "vulnerable-dependency"] + + +def test_duplicate_qualified_names_get_distinct_call_graph_names() -> None: + from coretrace_python.interprocedural import CallGraphAnalysis + from coretrace_python.taint import TaintAnalysis + + source = SourceManager().add_source( + "props.py", + "import os\n\nclass Box:\n @property\n def cmd(self):\n return self._cmd\n\n" + " @cmd.setter\n def cmd(self, value):\n self._cmd = value\n os.system(input())\n\n" + "def run():\n return 1\n\ndef run():\n return 2\n", + ) + module = engine.build_hir(source) + manager = engine.build_manager(module) + graph = manager.get(CallGraphAnalysis) + + assert graph.functions == ("Box.cmd", "Box.cmd__2", "run", "run__2") + setter = module.body[1].body[1] # type: ignore[union-attr] + assert graph.name_of(setter) == "Box.cmd__2" # type: ignore[arg-type] + assert manager.get(TaintAnalysis, setter) is not None # type: ignore[arg-type] + findings = engine.check(source, [Path(__file__).resolve().parent.parent / "plugins"]) + assert [(f.rule_id, f.span.start_line) for f in findings] == [("command-injection", 11)] diff --git a/tests/test_django_http_clients.py b/tests/test_django_http_clients.py index c4c706f..eee857a 100644 --- a/tests/test_django_http_clients.py +++ b/tests/test_django_http_clients.py @@ -91,7 +91,7 @@ def test_parameter_annotations_are_kept() -> None: def test_unsupported_annotations_do_not_break_the_build() -> None: - module = module_for("def f(x: (lambda: 1), y: 'Request'):\n return x\n") + module = module_for("def f(x: (n := 1), y: 'Request'):\n return x\n") function = module.body[0] assert isinstance(function, nodes.Function) assert function.parameters[0].annotation is None diff --git a/tests/test_interprocedural.py b/tests/test_interprocedural.py index 3fd2d5a..71ff83f 100644 --- a/tests/test_interprocedural.py +++ b/tests/test_interprocedural.py @@ -117,7 +117,7 @@ def test_methods_are_named_by_their_class() -> None: def test_recursion_and_unsupported_functions() -> None: graph = call_graph( "def loop(n):\n return loop(n)\n\n" - "def outer():\n def inner():\n pass\n return inner\n" + "def outer():\n class Inner:\n pass\n return Inner\n" ) assert graph.callees("loop") == frozenset({"loop"}) @@ -227,7 +227,7 @@ def test_recursive_summaries_reach_a_fixpoint() -> None: def test_unsupported_functions_have_conservative_summaries() -> None: - table = summaries("def outer(a):\n def inner():\n pass\n return inner\n\ndef use(b):\n return outer(b)\n") + table = summaries("def outer(a):\n class Inner:\n pass\n return Inner\n\ndef use(b):\n return outer(b)\n") assert table.summary("outer").return_dependencies == frozenset({0}) assert table.summary("outer").unsupported is True diff --git a/tests/test_parallel.py b/tests/test_parallel.py index 56d75ce..b3fe025 100644 --- a/tests/test_parallel.py +++ b/tests/test_parallel.py @@ -170,7 +170,7 @@ def test_parallel_analysis_keeps_syntax_and_unsupported_notes(tmp_path: Path) -> tmp_path / "src", { "broken.py": "def (:\n", - "nested.py": "def outer():\n def inner():\n pass\n return inner\n", + "nested.py": "def outer():\n class Inner:\n pass\n return Inner\n", "fine.py": CLEAN, }, ) diff --git a/tests/test_project.py b/tests/test_project.py index 742de43..8a5cab0 100644 --- a/tests/test_project.py +++ b/tests/test_project.py @@ -215,7 +215,7 @@ def test_unsupported_functions_and_syntax_errors_are_reported_per_file(tmp_path: tmp_path, { "good.py": "def f():\n pass\n", - "nested.py": "def outer():\n def inner():\n pass\n return inner\n", + "nested.py": "def outer():\n class Inner:\n pass\n return Inner\n", "broken.py": "def broken(:\n", }, ) diff --git a/tests/test_syntax_control_flow.py b/tests/test_syntax_control_flow.py new file mode 100644 index 0000000..6604287 --- /dev/null +++ b/tests/test_syntax_control_flow.py @@ -0,0 +1,269 @@ +"""Acceptance tests for the control-flow syntax found blocking real repositories +(``docs/architecture.md`` §3.2, §5, §6; second post-roadmap syntax pass). + +Conditional expressions and comprehensions carry control flow inside an expression. The +HIR keeps them as they are, ``Conditional`` and ``Comprehension``, so scopes stay +faithful; the CFG builder lays them out as real blocks, a branch or a loop writing a +synthetic local that the expression then reads, so PyIR, SSA, taint and refutation see +ordinary control flow. Lambdas and nested function definitions become ``MakeFunction`` +values whose bodies are not analysed yet; assignments to ``global`` names become +``SetGlobal``; set literals build ``BuildSet``; a chained assignment assigns each target +from one hidden local. + +Expected to remain red until ``Conditional``, ``Lambda``, ``Set``, dictionary and tuple +target comprehensions, ``MakeFunction``, ``BuildSet`` and ``SetGlobal`` exist. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from coretrace_python import engine +from coretrace_python.analysis import AnalysisManager +from coretrace_python.frontend import build_hir +from coretrace_python.hir import nodes +from coretrace_python.ir.lowering import analyzable_functions, lower_module +from coretrace_python.ir.model import Branch, ForNext, FunctionIR, Phi +from coretrace_python.ir.printer import format_module +from coretrace_python.semantic.scopes import BindingKind, ResolutionKind, ScopeKind, analyze_scopes +from coretrace_python.semantic.symbols import SymbolId +from coretrace_python.source import SourceManager +from coretrace_python.taint import ( + SecurityModelAnalysis, + SecurityModelRegistry, + Sink, + Source, + TaintAnalysis, + TaintKind, +) + +try: + from coretrace_python.hir.nodes import Conditional, Lambda, Set + from coretrace_python.ir.model import BuildSet, MakeFunction, SetGlobal +except ImportError as error: # pragma: no cover - red until the syntax lands + MISSING = error +else: + MISSING = None + + +@pytest.fixture(autouse=True) +def require_syntax() -> None: + if MISSING is not None: + pytest.fail(f"control-flow syntax pass is not implemented yet: {MISSING}") + + +REPO = Path(__file__).resolve().parent.parent +PLUGINS = REPO / "plugins" + +MODELS = ( + Source(SymbolId("python.builtins.input"), "stdin"), + Sink(SymbolId("python.os.system"), TaintKind.COMMAND), +) + + +def hir(text: str) -> nodes.Module: + return build_hir(SourceManager().add_source("s.py", text)) + + +def lower(text: str, *, ssa: bool = False) -> FunctionIR: + return lower_module(hir(text), ssa=ssa).functions[0] + + +def printed(text: str) -> str: + return format_module(lower_module(hir(text))) + + +def manager_for(text: str) -> AnalysisManager: + manager = AnalysisManager(hir("import os\n\n" + text)) + manager.register(*engine.ALL_ANALYSES) + registry = SecurityModelRegistry() + registry.register(*MODELS) + manager.provide(SecurityModelAnalysis, registry.freeze()) + return manager + + +def flow_lines(text: str, name: str | None = None) -> list[int]: + manager = manager_for(text) + functions = [s for s in manager.module.body if isinstance(s, nodes.Function)] + function = functions[0] if name is None else next(f for f in functions if f.name == name) + return sorted(f.location.start_line for f in manager.get(TaintAnalysis, function).flows) + + +def notes(text: str) -> list[str]: + findings = engine.check(SourceManager().add_source("m.py", text), [PLUGINS]) + return [f.message for f in findings if f.rule_id in ("syntax-error", "unsupported-syntax")] + + +def first_statement(text: str) -> nodes.Statement: + function = hir(text).body[0] + assert isinstance(function, nodes.Function) + return function.body[0] + + +# --------------------------------------------------------------------------- conditional expressions + + +def test_conditional_expressions_are_kept_in_the_hir_and_branch_in_the_cfg() -> None: + returned = first_statement("def f(a, b, c):\n return a if c else b\n") + assert isinstance(returned, nodes.Return) and isinstance(returned.value, Conditional) + assert isinstance(returned.value.test, nodes.Name) and returned.value.test.identifier == "c" + + function = lower("def f(a, b, c):\n return a if c else b\n", ssa=True) + + assert any(isinstance(b.terminator, Branch) for b in function.blocks) + assert any(isinstance(i, Phi) for b in function.blocks for i in b.instructions) + + +def test_taint_flows_through_either_branch_of_a_conditional() -> None: + assert flow_lines("def f(flag):\n cmd = 'ls' if flag else input()\n os.system(cmd)\n") == [5] + assert flow_lines("def f(flag):\n cmd = 'ls' if flag else 'pwd'\n os.system(cmd)\n") == [] + assert flow_lines("def f(flag):\n os.system(input() if flag else 'ls')\n") == [4] + + +def test_conditionals_nest_inside_other_expressions() -> None: + assert flow_lines("def f(a, b):\n os.system('ping ' + (input() if a else ('x' if b else 'y')))\n") == [4] + + +def test_conditionals_in_loop_conditions_are_reported_not_crashed() -> None: + found = notes("def f(x):\n while (x if x else 1):\n x = x - 1\n") + assert len(found) == 1 and "loop condition" in found[0] + + +# --------------------------------------------------------------------------- lambdas and nested functions + + +def test_lambdas_have_their_own_scope_and_lower_to_function_values() -> None: + module = hir("def f(items, k):\n return sorted(items, key=lambda x: x + k)\n") + scopes = analyze_scopes(module) + f = next(s for s in scopes.children(scopes.module_scope.id) if s.name == "f") + (lam,) = [s for s in scopes.children(f.id) if s.name == ""] + + assert lam.kind is ScopeKind.FUNCTION + assert lam.bindings["x"].kind is BindingKind.PARAMETER + assert scopes.resolve(lam.id, "k").kind is ResolutionKind.FREE + assert "x" not in f.bindings + returned = module.body[0].body[0] # type: ignore[union-attr] + assert isinstance(returned, nodes.Return) and isinstance(returned.value, nodes.Call) + assert isinstance(returned.value.keywords[0].value, Lambda) + + text = printed("def f(items, k):\n return sorted(items, key=lambda x: x + k)\n") + assert 'make_function ""' in text + (made,) = [i for b in lower("def f():\n return lambda: 1\n").blocks for i in b.instructions if isinstance(i, MakeFunction)] + assert made.name == "" + + +def test_calling_a_lambda_is_conservative() -> None: + assert flow_lines("def f():\n run = lambda c: c\n os.system(run(input()))\n") == [5] + + +def test_nested_functions_become_values_and_the_outer_function_is_analysed() -> None: + module = hir("def outer():\n def inner(x):\n return x\n os.system(inner(input()))\n") + assert [f.name for f in analyzable_functions(module)] == ["outer"] + + text = printed("def outer():\n def inner(x):\n return x\n return inner\n") + assert 'make_function "inner"' in text and 'store_local "inner"' in text + assert flow_lines("def outer():\n def inner(x):\n return x\n os.system(inner(input()))\n") == [6] + assert notes("def outer():\n def inner(x):\n return x\n return inner\n") == [] + + +# --------------------------------------------------------------------------- comprehensions + + +def test_comprehensions_lower_to_loops_over_a_synthetic_result() -> None: + function = lower("def f(items):\n return [x.strip() for x in items if x]\n") + text = printed("def f(items):\n return [x.strip() for x in items if x]\n") + + assert any(isinstance(b.terminator, ForNext) for b in function.blocks) + assert "append" in text and "build_list" in text + + +def test_dict_and_tuple_target_comprehensions_are_represented() -> None: + returned = first_statement("def f(pairs):\n return {k: v for k, v in pairs if k}\n") + assert isinstance(returned, nodes.Return) and isinstance(returned.value, nodes.Comprehension) + comprehension = returned.value + assert comprehension.kind == "dict" + assert isinstance(comprehension.key, nodes.Name) and comprehension.key.identifier == "k" + assert isinstance(comprehension.generators[0].target, nodes.Tuple) + scopes = analyze_scopes(hir("def f(pairs):\n return {k: v for k, v in pairs if k}\n")) + f = next(s for s in scopes.children(scopes.module_scope.id) if s.name == "f") + (inner,) = [s for s in scopes.children(f.id) if s.name == ""] + assert {"k", "v"} <= set(inner.bindings) + + +@pytest.mark.parametrize( + "body", + [ + "def f():\n cmds = [c for c in input().split()]\n os.system(cmds[0])\n", + "def f():\n cmds = {c for c in input().split()}\n os.system(list(cmds)[0])\n", + "def f():\n pairs = [('a', input())]\n d = {k: v for k, v in pairs}\n os.system(d['a'])\n", + "def f():\n os.system(' '.join(c for c in input().split()))\n", + "def f():\n xs = [[input()]]\n flat = [y for x in xs for y in x]\n os.system(flat[0])\n", + "def f():\n cmds = [c.strip() for c in input().split() if c]\n os.system(cmds[0])\n", + ], +) +def test_taint_flows_through_comprehensions(body: str) -> None: + found = flow_lines(body) + assert len(found) == 1 + assert found[0] == body.count("\n") + 2 + + +def test_comprehension_variables_do_not_leak_into_the_function() -> None: + assert flow_lines("def f():\n x = 'safe'\n ys = [x for x in input().split()]\n os.system(x)\n") == [] + assert flow_lines("def f():\n x = 'safe'\n ys = [x for x in input().split()]\n os.system(ys[0])\n") == [6] + + +def test_clean_comprehensions_stay_clean() -> None: + assert flow_lines("def f(items):\n names = [i.name for i in items]\n os.system(names[0])\n") == [] + + +# --------------------------------------------------------------------------- globals, sets, chains + + +def test_assignments_to_global_names_lower_to_set_global() -> None: + text = printed("counter = 0\n\ndef f():\n global counter\n counter = 1\n counter += 1\n") + assert text.count('set_global "counter"') == 2 + (store, _) = [i for b in lower("counter = 0\n\ndef f():\n global counter\n counter = 1\n counter += 1\n").blocks for i in b.instructions if isinstance(i, SetGlobal)] + assert store.name == "counter" + assert notes("counter = 0\n\ndef f():\n global counter\n counter = input()\n") == [] + + +def test_set_literals_build_sets_and_carry_taint() -> None: + returned = first_statement("def f():\n return {1, 2}\n") + assert isinstance(returned, nodes.Return) and isinstance(returned.value, Set) + (built,) = [i for b in lower("def f():\n return {1, 2}\n").blocks for i in b.instructions if isinstance(i, BuildSet)] + assert len(built.elements) == 2 + assert "build_set" in printed("def f():\n return {1, 2}\n") + assert flow_lines("def f():\n s = {input()}\n os.system(list(s)[0])\n") == [5] + + +def test_chained_assignments_assign_every_target() -> None: + function = hir("def f():\n a = b = input()\n return a, b\n").body[0] + assert isinstance(function, nodes.Function) + assigns = [s for s in function.body if isinstance(s, nodes.Assign)] + assert len(assigns) == 3 + assert flow_lines("def f():\n a = b = input()\n os.system(a)\n os.system(b)\n") == [5, 6] + + +# --------------------------------------------------------------------------- end to end + + +def test_a_realistic_module_is_fully_analysed() -> None: + assert ( + notes( + "import os\nimport json\n\n" + "state = {}\n\n" + "def configure(settings):\n" + " global state\n" + " state = {k: v for k, v in settings.items() if not k.startswith('_')}\n" + " ordered = sorted(state, key=lambda k: k.lower())\n" + " label = ordered[0] if ordered else 'none'\n" + " tags = {t.strip() for t in label.split(',')}\n" + " first = last = label\n" + " def render(item):\n" + " return json.dumps(item)\n" + " return render(tags), first, last\n" + ) + == [] + ) diff --git a/tests/test_syntax_coverage.py b/tests/test_syntax_coverage.py index d9589ac..442394a 100644 --- a/tests/test_syntax_coverage.py +++ b/tests/test_syntax_coverage.py @@ -266,9 +266,9 @@ def test_check_reports_unsupported_functions_and_keeps_going(tmp_path, capsys) - source.write_text( "CONFIG = {}\n\n" "def outer():\n" - " def inner():\n" - " return 1\n" - " return inner\n\n" + " class Inner:\n" + " pass\n" + " return Inner\n\n" "class Runner:\n" " def go(self, code):\n" " eval(code)\n", @@ -280,13 +280,13 @@ def test_check_reports_unsupported_functions_and_keeps_going(tmp_path, capsys) - assert exit_code == 1 assert f"{source}:3:1: info unsupported-syntax: " in output - assert "Function" in output + assert "Class" in output assert f"{source}:10:9: high dangerous-eval:" in output assert output.endswith("2 findings\n") def test_unsupported_syntax_findings_are_notes() -> None: - findings = engine.check(SourceManager().add_source("n.py", "def f():\n def g():\n pass\n"), [PLUGINS]) + findings = engine.check(SourceManager().add_source("n.py", "def f():\n class G:\n pass\n"), [PLUGINS]) assert [f.rule_id for f in findings] == ["unsupported-syntax"] assert findings[0].severity is Severity.INFO From 316652c900e5c1c0036abd7f218f12ac562cb9f3 Mon Sep 17 00:00:00 2001 From: Hugo Date: Fri, 4 Sep 2026 06:18:42 +0200 Subject: [PATCH 2/2] feat: lay out conditional expressions and comprehensions in the CFG, lower lambdas, nested functions, sets and global stores --- README.md | 9 +- src/coretrace_python/abstract/heap.py | 12 +- src/coretrace_python/cfg/builder.py | 196 +++++++++++++++++- src/coretrace_python/cfg/model.py | 3 + src/coretrace_python/frontend/ast_adapter.py | 39 +++- src/coretrace_python/hir/nodes.py | 34 ++- .../interprocedural/callgraph.py | 14 +- .../interprocedural/summaries.py | 3 +- src/coretrace_python/ir/lowering.py | 33 ++- src/coretrace_python/ir/model.py | 31 ++- src/coretrace_python/ir/printer.py | 10 + src/coretrace_python/semantic/scopes.py | 21 +- src/coretrace_python/taint/engine.py | 3 +- 13 files changed, 374 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index b353594..86159c8 100644 --- a/README.md +++ b/README.md @@ -140,9 +140,12 @@ forms, decorators, assignments to names, attributes, items and unpacked tuples, assignment, list, tuple and dict literals with `*` and `**` unpacking, f-strings, slices, `and`/`or`, chained comparisons, keyword and starred arguments, `with`, `assert`, `try`/`except`/`else`/`finally`, `await`, `yield`, `if`/`elif`/`else`, `while`, `for` -with name or tuple targets, `break`, `continue` and `raise`, `from` clause included. Not -yet: conditional expressions, lambdas, comprehensions, nested function definitions and -assignments to `global` names. Blocks inside a `try` +with name or tuple targets, `break`, `continue` and `raise`, `from` clause included, +conditional expressions and comprehensions, laid out as real branches and loops over a +synthetic local, set literals, chained assignments, assignments to `global` names, and +lambdas and nested function definitions as function values whose bodies are not analysed +yet. Not yet: `match`, `del`, `nonlocal` assignments, classes defined inside functions, and +a conditional expression or comprehension inside a `while` condition. Blocks inside a `try` body carry exception edges to the handlers; `finally` is modelled on the normal path only. Methods of module-level classes are analysed like functions; other module-level code is skipped. An import inside a function is shown where it runs, as an `import` instruction naming the module, the bound name and the canonical symbol. A function using syntax outside this subset is reported by `--check` as diff --git a/src/coretrace_python/abstract/heap.py b/src/coretrace_python/abstract/heap.py index e6b52ab..7baf8e0 100644 --- a/src/coretrace_python/abstract/heap.py +++ b/src/coretrace_python/abstract/heap.py @@ -23,6 +23,7 @@ Await, BuildDict, BuildList, + BuildSet, BuildTuple, Call, Catch, @@ -183,11 +184,12 @@ def instruction(self, instruction: Instruction) -> bool: result = instruction.result if result is None: return False - if isinstance(instruction, BuildList | BuildTuple | BuildDict | Call | WithEnter | Catch | Yield): + if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet | Call | WithEnter | Catch | Yield): kind = { BuildList: "list", BuildTuple: "tuple", BuildDict: "dict", + BuildSet: "set", Call: "call", WithEnter: "context", Catch: "exception", @@ -195,11 +197,11 @@ def instruction(self, instruction: Instruction) -> bool: }[type(instruction)] site = AbstractObject(AllocationSite(kind, instruction.location)) changed = self.assign(result, {site}) - if isinstance(instruction, BuildList | BuildTuple | BuildDict): + if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet): values = ( - instruction.elements - if isinstance(instruction, BuildList | BuildTuple) - else tuple(v for _, v in instruction.items) + tuple(v for _, v in instruction.items) + if isinstance(instruction, BuildDict) + else instruction.elements ) for element in values: changed |= self.store({site}, ELEMENTS, element) diff --git a/src/coretrace_python/cfg/builder.py b/src/coretrace_python/cfg/builder.py index 674bb1a..0db932c 100644 --- a/src/coretrace_python/cfg/builder.py +++ b/src/coretrace_python/cfg/builder.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import ClassVar +from dataclasses import fields, is_dataclass, replace +from typing import Any, ClassVar from coretrace_python.analysis import AnalysisContext, FunctionAnalysis from coretrace_python.cfg.model import ( @@ -36,13 +37,137 @@ def __init__(self, function: nodes.Function) -> None: self.counters: dict[str, int] = {} self.loops: list[tuple[BlockId, BlockId]] = [] self.handlers: list[tuple[BlockId, ...]] = [] + self.synthetic: set[str] = set() def build(self) -> CFG: entry = BlockId("entry") end = self.sequence(self.function.body, _Open(entry), None, self.function.span) if end is not None: self.finish(end, Return(None, self.function.span)) - return CFG(entry, self.blocks) + return CFG(entry, self.blocks, frozenset(self.synthetic)) + + # ------------------------------------------------------------------ expression-level control flow + + def hidden(self, kind: str, span: SourceSpan) -> nodes.Name: + name = f"_coretrace_{kind}_{span.start_line}_{span.start_column}_{len(self.synthetic)}" + self.synthetic.add(name) + return nodes.Name(name, span) + + def desugared(self, node: nodes.Statement, block: _Open) -> tuple[nodes.Statement, _Open]: + """``node`` with its conditional expressions and comprehensions replaced by reads + of synthetic locals, after laying out the statements that compute them.""" + + pending: list[nodes.Statement] = [] + if isinstance(node, nodes.While): + if _has_control_flow(node.condition): + raise CFGError( + f"{node.span.display()}: conditional expressions and comprehensions in a " + "loop condition are not supported yet" + ) + return node, block + if isinstance(node, nodes.Assign | nodes.AugAssign): + node = replace(node, target=self.hoist(node.target, pending), value=self.hoist(node.value, pending)) + elif isinstance(node, nodes.ExpressionStatement): + node = replace(node, expression=self.hoist(node.expression, pending)) + elif isinstance(node, nodes.Return | nodes.Raise) and node.__class__ is nodes.Return: + if node.value is not None: + node = replace(node, value=self.hoist(node.value, pending)) + elif isinstance(node, nodes.Raise): + exception = self.hoist(node.exception, pending) if node.exception is not None else None + cause = self.hoist(node.cause, pending) if node.cause is not None else None + node = replace(node, exception=exception, cause=cause) + elif isinstance(node, nodes.Assert): + message = self.hoist(node.message, pending) if node.message is not None else None + node = replace(node, test=self.hoist(node.test, pending), message=message) + elif isinstance(node, nodes.If): + node = replace(node, condition=self.hoist(node.condition, pending)) + elif isinstance(node, nodes.For): + node = replace(node, iterable=self.hoist(node.iterable, pending)) + elif isinstance(node, nodes.With): + items = tuple(replace(item, context=self.hoist(item.context, pending)) for item in node.items) + node = replace(node, items=items) + if not pending: + return node, block + laid_out = self.sequence(tuple(pending), block, None, node.span) + assert laid_out is not None, "hoisted statements always fall through" + return node, laid_out + + def hoist(self, node: Any, pending: list[nodes.Statement]) -> Any: + """Rewrite one expression tree, appending the statements it needs to ``pending``.""" + + if isinstance(node, nodes.Conditional): + test = self.hoist(node.test, pending) + result = self.hidden("cond", node.span) + pending.append( + nodes.If( + test, + (nodes.Assign(result, node.body, node.span),), + (nodes.Assign(result, node.orelse, node.span),), + node.span, + ) + ) + return result + if isinstance(node, nodes.Comprehension): + return self.comprehension(node, pending) + if isinstance(node, nodes.Lambda): + defaults = tuple( + replace(p, default=self.hoist(p.default, pending)) if p.default is not None else p + for p in node.parameters + ) + return replace(node, parameters=defaults) + if isinstance(node, tuple): + return tuple(self.hoist(item, pending) for item in node) + if is_dataclass(node) and not isinstance(node, type): + changes = { + f.name: self.hoist(getattr(node, f.name), pending) + for f in fields(node) + if f.name != "span" and _may_hold_expressions(getattr(node, f.name)) + } + return replace(node, **changes) if changes else node + return node + + def comprehension(self, node: nodes.Comprehension, pending: list[nodes.Statement]) -> nodes.Name: + """Lay a comprehension out as loops filling a synthetic collection.""" + + span = node.span + result = self.hidden("comp", span) + renames: dict[str, str] = {} + for generator in node.generators: + for name in _bound_names(generator.target): + renames[name] = self.hidden(f"var_{name}", generator.target.span).identifier + if node.kind == "set": + initial: nodes.Expression = nodes.Call(nodes.Name("set", span), (), (), span) + elif node.kind == "dict": + initial = nodes.Dict((), span) + else: + initial = nodes.List((), span) + pending.append(nodes.Assign(result, initial, span)) + + element = _rename(node.element, renames) + if node.kind == "dict": + assert node.key is not None + innermost: nodes.Statement = nodes.Assign( + nodes.Subscript(result, _rename(node.key, renames), span), element, span + ) + else: + method = "add" if node.kind == "set" else "append" + call = nodes.Call(nodes.Attribute(result, method, span), (element,), (), span) + innermost = nodes.ExpressionStatement(call, span) + + body: tuple[nodes.Statement, ...] = (innermost,) + for index in range(len(node.generators) - 1, -1, -1): + generator = node.generators[index] + for condition in reversed(generator.conditions): + body = (nodes.If(_rename(condition, renames), body, (), condition.span),) + iterable = _rename(generator.iterable, renames) if index else self.hoist(generator.iterable, pending) + if isinstance(generator.target, nodes.Name): + target = nodes.Name(renames[generator.target.identifier], generator.target.span) + else: + target = self.hidden("item", generator.target.span) + body = (nodes.Assign(_rename_target(generator.target, renames), target, generator.target.span), *body) + body = (nodes.For(target, iterable, body, False, generator.span),) + pending.extend(body) + return result # ------------------------------------------------------------------ blocks @@ -92,6 +217,7 @@ def sequence( def statement( self, node: nodes.Statement, block: _Open, continuation: BlockId | None ) -> _Open | None: + node, block = self.desugared(node, block) if isinstance(node, nodes.Return): self.finish(block, Return(node.value, node.span)) return None @@ -201,6 +327,72 @@ def loop(self, keyword: str, span: SourceSpan) -> tuple[BlockId, BlockId]: return self.loops[-1] +def _may_hold_expressions(value: object) -> bool: + return isinstance(value, tuple) or (is_dataclass(value) and not isinstance(value, type)) + + +def _has_control_flow(node: object) -> bool: + if isinstance(node, nodes.Conditional | nodes.Comprehension): + return True + if isinstance(node, nodes.Lambda): + return False + if isinstance(node, tuple): + return any(_has_control_flow(item) for item in node) + if is_dataclass(node) and not isinstance(node, type): + return any(_has_control_flow(getattr(node, f.name)) for f in fields(node) if f.name != "span") + return False + + +def _bound_names(target: nodes.Target) -> list[str]: + if isinstance(target, nodes.Name): + return [target.identifier] + if isinstance(target, nodes.Tuple): + return [name for element in target.elements for name in _bound_names(element)] # type: ignore[arg-type] + return [] + + +def _rename(node: Any, renames: dict[str, str]) -> Any: + """``node`` with comprehension variables replaced by their synthetic locals, without + entering scopes that rebind them.""" + + if not renames: + return node + if isinstance(node, nodes.Name): + return nodes.Name(renames[node.identifier], node.span) if node.identifier in renames else node + if isinstance(node, nodes.Lambda): + inner = {k: v for k, v in renames.items() if k not in {p.name for p in node.parameters}} + return replace(node, body=_rename(node.body, inner)) + if isinstance(node, nodes.Comprehension): + shadowed = {name for g in node.generators for name in _bound_names(g.target)} + inner = {k: v for k, v in renames.items() if k not in shadowed} + generators = tuple( + replace( + g, + iterable=_rename(g.iterable, renames if index == 0 else inner), + conditions=_rename(g.conditions, inner), + ) + for index, g in enumerate(node.generators) + ) + key = _rename(node.key, inner) if node.key is not None else None + return replace(node, element=_rename(node.element, inner), generators=generators, key=key) + if isinstance(node, tuple): + return tuple(_rename(item, renames) for item in node) + if is_dataclass(node) and not isinstance(node, type): + changes = { + f.name: _rename(getattr(node, f.name), renames) + for f in fields(node) + if f.name != "span" and _may_hold_expressions(getattr(node, f.name)) + } + return replace(node, **changes) if changes else node + return node + + +def _rename_target(target: nodes.Target, renames: dict[str, str]) -> nodes.Target: + renamed = _rename(target, renames) + assert isinstance(renamed, nodes.Name | nodes.Tuple | nodes.Attribute | nodes.Subscript) + return renamed + + def build_cfg(function: nodes.Function) -> CFG: return _Builder(function).build() diff --git a/src/coretrace_python/cfg/model.py b/src/coretrace_python/cfg/model.py index 8d61123..2f57779 100644 --- a/src/coretrace_python/cfg/model.py +++ b/src/coretrace_python/cfg/model.py @@ -92,6 +92,9 @@ class BasicBlock: class CFG: entry: BlockId blocks: Mapping[BlockId, BasicBlock] + # Locals the builder introduced while laying out expression-level control flow + # (conditional expressions, comprehensions); the lowering treats them as locals. + synthetic_locals: frozenset[str] = frozenset() _successors: Mapping[BlockId, tuple[BlockId, ...]] = field( init=False, repr=False, compare=False ) diff --git a/src/coretrace_python/frontend/ast_adapter.py b/src/coretrace_python/frontend/ast_adapter.py index 3c81e17..4d97baa 100644 --- a/src/coretrace_python/frontend/ast_adapter.py +++ b/src/coretrace_python/frontend/ast_adapter.py @@ -29,7 +29,7 @@ class HIRBuildError(Exception): ast.MatMult: "matmul", } _UNARY_OPERATORS = {ast.Invert: "invert", ast.Not: "not", ast.UAdd: "pos", ast.USub: "neg"} -_COMPREHENSIONS = {ast.ListComp: "list", ast.SetComp: "set", ast.GeneratorExp: "generator"} +_COMPREHENSIONS = {ast.ListComp: "list", ast.SetComp: "set", ast.GeneratorExp: "generator", ast.DictComp: "dict"} _COMPARE_OPERATORS = { ast.Eq: "eq", ast.NotEq: "not_eq", @@ -118,6 +118,17 @@ def expression(self, node: ast.expr) -> nodes.Expression: ) if isinstance(node, ast.Starred): return nodes.Starred(self.expression(node.value), span) + if isinstance(node, ast.Set): + return nodes.Set(self.elements(node.elts), span) + if isinstance(node, ast.IfExp): + return nodes.Conditional( + self.expression(node.test), self.expression(node.body), self.expression(node.orelse), span + ) + if isinstance(node, ast.Lambda): + return nodes.Lambda(self.parameters(node.args), self.expression(node.body), span) + if isinstance(node, ast.DictComp): + generators = tuple(self.generator(generator) for generator in node.generators) + return nodes.Comprehension("dict", self.expression(node.value), generators, span, self.expression(node.key)) if isinstance(node, ast.Attribute): return nodes.Attribute(self.expression(node.value), node.attr, span) if isinstance(node, ast.Subscript): @@ -176,9 +187,7 @@ def target(self, node: ast.expr) -> nodes.Target: def generator(self, node: ast.comprehension) -> nodes.ComprehensionGenerator: if node.is_async: self.fail(node.target, "async comprehensions are not supported yet") - if not isinstance(node.target, ast.Name): - self.fail(node.target, "only a single name is supported as a comprehension target") - target = nodes.Name(node.target.id, self.span(node.target)) + target = self.target(node.target) conditions = tuple(self.expression(condition) for condition in node.ifs) # ``ast.comprehension`` carries no location; span it from the target to the last clause. last = self.span(node.ifs[-1] if node.ifs else node.iter) @@ -191,11 +200,21 @@ def generator(self, node: ast.comprehension) -> nodes.ComprehensionGenerator: ) return nodes.ComprehensionGenerator(target, self.expression(node.iter), conditions, span) + def statements(self, node: ast.stmt) -> list[nodes.Statement]: + """The HIR statements of one source statement: several for ``a = b = value``, + which assigns a hidden local once and every target from it.""" + + span = self.span(node) + if isinstance(node, ast.Assign) and len(node.targets) > 1: + hidden = nodes.Name(f"_coretrace_chain_{span.start_line}_{span.start_column}", span) + found: list[nodes.Statement] = [nodes.Assign(hidden, self.expression(node.value), span)] + found.extend(nodes.Assign(self.target(target), hidden, self.span(target)) for target in node.targets) + return found + return [self.statement(node)] + def statement(self, node: ast.stmt) -> nodes.Statement: span = self.span(node) if isinstance(node, ast.Assign): - if len(node.targets) != 1: - self.fail(node, "chained assignment is not supported yet") return nodes.Assign(self.target(node.targets[0]), self.expression(node.value), span) if isinstance(node, ast.AnnAssign): # The annotation never affects behaviour; a bare declaration does nothing. @@ -309,18 +328,18 @@ def statement(self, node: ast.stmt) -> nodes.Statement: self.fail(node) def block(self, statements: list[ast.stmt]) -> tuple[nodes.Statement, ...]: - return tuple(self.statement(statement) for statement in statements) + return tuple(found for statement in statements for found in self.statements(statement)) def class_definition(self, node: ast.ClassDef) -> nodes.Class: if node.keywords: self.fail(node, "class keyword arguments are not supported yet") bases = tuple(self.expression(base) for base in node.bases) - body = tuple(self.statement(statement) for statement in node.body) + body = self.block(node.body) decorators = tuple(self.expression(d) for d in node.decorator_list) return nodes.Class(node.name, bases, body, self.span(node), decorators) def function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> nodes.Function: - body = tuple(self.statement(statement) for statement in node.body) + body = self.block(node.body) return nodes.Function( name=node.name, parameters=self.parameters(node.args), @@ -357,7 +376,7 @@ def parameter(self, argument: ast.arg, default: ast.expr | None, kind: str) -> n return nodes.Parameter(argument.arg, self.span(argument), value, kind, annotation) def module(self, tree: ast.Module) -> nodes.Module: - body = tuple(self.statement(statement) for statement in tree.body) + body = self.block(tree.body) if tree.body: first_span = self.span(tree.body[0]) last_span = self.span(tree.body[-1]) diff --git a/src/coretrace_python/hir/nodes.py b/src/coretrace_python/hir/nodes.py index 51a32aa..3171748 100644 --- a/src/coretrace_python/hir/nodes.py +++ b/src/coretrace_python/hir/nodes.py @@ -130,6 +130,31 @@ class Starred: span: SourceSpan +@dataclass(frozen=True, slots=True) +class Set: + elements: tuple[Expression, ...] + span: SourceSpan + + +@dataclass(frozen=True, slots=True) +class Conditional: + """``body if test else orelse``; the CFG builder lays it out as a branch.""" + + test: Expression + body: Expression + orelse: Expression + span: SourceSpan + + +@dataclass(frozen=True, slots=True) +class Lambda: + """An anonymous function; its body is a scope of its own and is not lowered.""" + + parameters: tuple[Parameter, ...] + body: Expression + span: SourceSpan + + @dataclass(frozen=True, slots=True) class Await: value: Expression @@ -146,7 +171,7 @@ class Yield: class ComprehensionGenerator: """One ``for target in iterable if condition...`` clause of a comprehension.""" - target: Name + target: Target iterable: Expression conditions: tuple[Expression, ...] span: SourceSpan @@ -154,12 +179,14 @@ class ComprehensionGenerator: @dataclass(frozen=True, slots=True) class Comprehension: - """A list, set or generator comprehension; ``kind`` names which one.""" + """A list, set, dict or generator comprehension; ``kind`` names which one. A dict + comprehension's ``element`` is the value and ``key`` its key.""" kind: str element: Expression generators: tuple[ComprehensionGenerator, ...] span: SourceSpan + key: Expression | None = None Expression: TypeAlias = ( @@ -178,6 +205,9 @@ class Comprehension: | FormattedString | Slice | Starred + | Set + | Conditional + | Lambda | Await | Yield | Comprehension diff --git a/src/coretrace_python/interprocedural/callgraph.py b/src/coretrace_python/interprocedural/callgraph.py index 1ec5d23..53429fb 100644 --- a/src/coretrace_python/interprocedural/callgraph.py +++ b/src/coretrace_python/interprocedural/callgraph.py @@ -158,10 +158,16 @@ class CallGraphAnalysis(Analysis[CallGraph]): @classmethod def compute(cls, ctx: AnalysisContext) -> CallGraph: scopes = ctx.get(ScopeAnalysis) - definitions = { - qualified_name(scopes, function): function - for function in analyzable_functions(ctx.module) - } + definitions: dict[str, nodes.Function] = {} + for function in analyzable_functions(ctx.module): + # A property and its setter, or a redefinition, share a qualified name; each + # definition still needs a name of its own. + base = unique = qualified_name(scopes, function) + ordinal = 2 + while unique in definitions: + unique = f"{base}__{ordinal}" + ordinal += 1 + definitions[unique] = function known = frozenset( name for name, function in definitions.items() if function in ctx.module.body ) diff --git a/src/coretrace_python/interprocedural/summaries.py b/src/coretrace_python/interprocedural/summaries.py index d4b7a56..f052e56 100644 --- a/src/coretrace_python/interprocedural/summaries.py +++ b/src/coretrace_python/interprocedural/summaries.py @@ -37,6 +37,7 @@ Branch, BuildDict, BuildList, + BuildSet, BuildTuple, Call, Constant, @@ -264,7 +265,7 @@ def instruction(self, instruction: Instruction, state: Mapping[Key, Dep]) -> Dep deps = EMPTY for operand in instruction.operands(): deps |= state.get(operand, EMPTY) - if isinstance(instruction, BuildList | BuildTuple | BuildDict): + if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet): for unpacked in instruction.unpacked: deps |= self.deep(unpacked, state) return deps diff --git a/src/coretrace_python/ir/lowering.py b/src/coretrace_python/ir/lowering.py index 23256cc..fe51b5e 100644 --- a/src/coretrace_python/ir/lowering.py +++ b/src/coretrace_python/ir/lowering.py @@ -25,6 +25,7 @@ Branch, BuildDict, BuildList, + BuildSet, BuildSlice, BuildString, BuildTuple, @@ -43,10 +44,12 @@ Instruction, Jump, LoadLocal, + MakeFunction, ModuleIR, Raise, Return, SetAttr, + SetGlobal, SetItem, StoreLocal, Symbol, @@ -108,6 +111,8 @@ def emit_effect(self, instruction: EffectInstruction) -> None: self.instructions.append(instruction) def resolve(self, name: str) -> Resolution: + if name in self.cfg.synthetic_locals: + return Resolution(ResolutionKind.LOCAL, self.scope.id) return self.scopes.resolve(self.scope.id, name) # ------------------------------------------------------------------ expressions @@ -170,6 +175,13 @@ def expression(self, node: nodes.Expression) -> Value: return self.emit(BuildSlice(self.new_value(), node.span, bounds[0], bounds[1], bounds[2])) if isinstance(node, nodes.Starred): self.fail(node, "a starred expression is only supported in calls, lists and tuples") + if isinstance(node, nodes.Set): + elements, unpacked = self.spread(node.elements) + return self.emit(BuildSet(self.new_value(), node.span, elements, unpacked)) + if isinstance(node, nodes.Lambda): + return self.emit(MakeFunction(self.new_value(), node.span, "")) + if isinstance(node, nodes.Conditional | nodes.Comprehension): + self.fail(node, "expression-level control flow must be laid out by the CFG builder") if isinstance(node, nodes.Await): return self.emit(Await(self.new_value(), node.span, self.expression(node.value))) if isinstance(node, nodes.Yield): @@ -200,8 +212,12 @@ def spread(self, elements: tuple[nodes.Expression, ...]) -> tuple[tuple[Value, . def store(self, target: nodes.Target, value: Value) -> None: if isinstance(target, nodes.Name): - if self.resolve(target.identifier).kind is not ResolutionKind.LOCAL: - self.fail(target, "assignment to a global or nonlocal name is not supported yet") + resolution = self.resolve(target.identifier) + if resolution.kind is ResolutionKind.GLOBAL: + self.emit_effect(SetGlobal(None, target.span, target.identifier, value)) + return + if resolution.kind is not ResolutionKind.LOCAL: + self.fail(target, "assignment to a nonlocal name is not supported yet") self.emit_effect(StoreLocal(None, target.span, target.identifier, value)) elif isinstance(target, nodes.Attribute): object_value = self.expression(target.value) @@ -221,6 +237,17 @@ def statement(self, node: nodes.Statement) -> None: if isinstance(node, nodes.Assign): self.store(node.target, self.expression(node.value)) return + if isinstance(node, nodes.Function): + # A nested definition is a value bound to its name; its body is its own + # scope and is not lowered here. + for decorator in node.decorators: + self.expression(decorator) + for parameter in node.parameters: + if parameter.default is not None: + self.expression(parameter.default) + made = self.emit(MakeFunction(self.new_value(), node.span, node.name)) + self.store(nodes.Name(node.name, node.span), made) + return if isinstance(node, nodes.AugAssign): current = self.expression(node.target) operand = self.expression(node.value) @@ -365,7 +392,7 @@ def walk(node: Node) -> None: names(node.target) elif isinstance(node, nodes.ExceptHandler) and node.name is not None: assigned.add(node.name) - if isinstance(node, nodes.Function | nodes.Class | nodes.Comprehension): + if isinstance(node, nodes.Function | nodes.Class | nodes.Comprehension | nodes.Lambda): return for child in children(node): walk(child) diff --git a/src/coretrace_python/ir/model.py b/src/coretrace_python/ir/model.py index f212003..de8c44a 100644 --- a/src/coretrace_python/ir/model.py +++ b/src/coretrace_python/ir/model.py @@ -156,6 +156,23 @@ class BuildDict(Operands): unpacked: tuple[Value, ...] = () +@dataclass(frozen=True) +class BuildSet(Operands): + result: Value + location: SourceSpan + elements: tuple[Value, ...] + unpacked: tuple[Value, ...] = () + + +@dataclass(frozen=True) +class MakeFunction(Operands): + """A function value: a lambda or a nested definition, whose body is not lowered.""" + + result: Value + location: SourceSpan + name: str + + @dataclass(frozen=True) class BuildString(Operands): """An f-string: its constant and formatted parts, in order.""" @@ -294,6 +311,16 @@ class Assert(Operands): message: Value | None +@dataclass(frozen=True) +class SetGlobal(Operands): + """Assignment to a name declared ``global``.""" + + result: None + location: SourceSpan + name: str + value: Value + + @dataclass(frozen=True) class Import(Operands): """An import executed where it stands: ``module`` as written (relative dots @@ -324,14 +351,16 @@ class Import(Operands): | BuildList | BuildTuple | BuildDict + | BuildSet | BuildString | BuildSlice + | MakeFunction | WithEnter | Catch | Await | Yield ) -EffectInstruction: TypeAlias = StoreLocal | SetAttr | SetItem | WithExit | Assert | Import +EffectInstruction: TypeAlias = StoreLocal | SetAttr | SetItem | WithExit | Assert | Import | SetGlobal Instruction: TypeAlias = ValueInstruction | EffectInstruction diff --git a/src/coretrace_python/ir/printer.py b/src/coretrace_python/ir/printer.py index b3e0f45..6aa7b25 100644 --- a/src/coretrace_python/ir/printer.py +++ b/src/coretrace_python/ir/printer.py @@ -8,6 +8,7 @@ Branch, BuildDict, BuildList, + BuildSet, BuildSlice, BuildString, BuildTuple, @@ -24,11 +25,13 @@ Instruction, Jump, LoadLocal, + MakeFunction, ModuleIR, Phi, Raise, Return, SetAttr, + SetGlobal, SetItem, StoreLocal, Symbol, @@ -85,6 +88,13 @@ def _instruction(instruction: Instruction) -> str: parts = [f"{_value(k)}: {_value(v)}" for k, v in instruction.items] parts.extend(f"**{_value(v)}" for v in instruction.unpacked) return f"{_value(instruction.result)} = build_dict {', '.join(parts)}".rstrip() + if isinstance(instruction, BuildSet): + parts = [_value(v) for v in instruction.elements] + [f"*{_value(v)}" for v in instruction.unpacked] + return f"{_value(instruction.result)} = build_set {', '.join(parts)}".rstrip() + if isinstance(instruction, MakeFunction): + return f'{_value(instruction.result)} = make_function "{instruction.name}"' + if isinstance(instruction, SetGlobal): + return f'set_global "{instruction.name}", {_value(instruction.value)}' if isinstance(instruction, BuildString): return f"{_value(instruction.result)} = build_string {', '.join(_value(v) for v in instruction.parts)}".rstrip() if isinstance(instruction, BuildSlice): diff --git a/src/coretrace_python/semantic/scopes.py b/src/coretrace_python/semantic/scopes.py index 907ed8a..8afda18 100644 --- a/src/coretrace_python/semantic/scopes.py +++ b/src/coretrace_python/semantic/scopes.py @@ -209,7 +209,7 @@ def freeze(self) -> Scope: ) -_COMPREHENSION_NAMES = {"list": "", "set": "", "generator": ""} +_COMPREHENSION_NAMES = {"list": "", "set": "", "generator": "", "dict": ""} class _Collector: @@ -358,6 +358,21 @@ def expression(self, node: nodes.Expression, scope: _ScopeBuilder) -> None: self.expression(bound, scope) elif isinstance(node, nodes.Starred): self.expression(node.value, scope) + elif isinstance(node, nodes.Set): + for element in node.elements: + self.expression(element, scope) + elif isinstance(node, nodes.Conditional): + self.expression(node.test, scope) + self.expression(node.body, scope) + self.expression(node.orelse, scope) + elif isinstance(node, nodes.Lambda): + for parameter in node.parameters: + if parameter.default is not None: + self.expression(parameter.default, scope) + inner = self.open(scope, ScopeKind.FUNCTION, "", node.span) + for parameter in node.parameters: + inner.bind(parameter.name, BindingKind.PARAMETER, parameter.span) + self.expression(node.body, inner) else: raise TypeError(f"unknown expression: {node!r}") @@ -379,9 +394,11 @@ def comprehension(self, node: nodes.Comprehension, scope: _ScopeBuilder) -> None for index, generator in enumerate(node.generators): if index: self.expression(generator.iterable, inner) - inner.bind(generator.target.identifier, BindingKind.LOCAL, generator.target.span) + self.target(generator.target, inner) for condition in generator.conditions: self.expression(condition, inner) + if node.key is not None: + self.expression(node.key, inner) self.expression(node.element, inner) def open( diff --git a/src/coretrace_python/taint/engine.py b/src/coretrace_python/taint/engine.py index 26132ce..790d273 100644 --- a/src/coretrace_python/taint/engine.py +++ b/src/coretrace_python/taint/engine.py @@ -42,6 +42,7 @@ Branch, BuildDict, BuildList, + BuildSet, BuildTuple, Call, Compare, @@ -254,7 +255,7 @@ def instruction(self, instruction: Instruction, state: Mapping[Key, Taint]) -> T return taint for operand in instruction.operands(): taint = taint.join(state.get(operand, Taint.none())) - if isinstance(instruction, BuildList | BuildTuple | BuildDict): + if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet): # ``[*xs]`` and ``{**d}`` copy the contents of what they unpack. for unpacked in instruction.unpacked: taint = taint.join(self.deep(unpacked, state))