diff --git a/src/renaissance/utils/cst_utils.py b/src/renaissance/utils/cst_utils.py new file mode 100644 index 00000000..4717f254 --- /dev/null +++ b/src/renaissance/utils/cst_utils.py @@ -0,0 +1,67 @@ +# Copyright (C) 2024-2026 TNO +# +# Licensed under the terms specified in the LICENSE file in the root directory +# of this repository, also available at: +# https://github.com/TNO/Renaissance.Py/blob/main/LICENSE + +"""Helpers over the libcst tree of one module: its top-level statements, the names they read or bind and the comment trailing a line.""" + +from collections.abc import Iterator + +import libcst as cst + + +def module_statements(module: cst.Module) -> Iterator[cst.BaseSmallStatement | cst.BaseCompoundStatement]: + """Yield the module's top-level statements in order, one per small statement of a line like `a = 1; b = 2`.""" + for statement in module.body: + if isinstance(statement, cst.SimpleStatementLine): + yield from statement.body + else: + yield statement + + +def trailing_of_sole_statement(module: cst.Module, statement: cst.BaseSmallStatement) -> cst.TrailingWhitespace | None: + """Return the whitespace and comment trailing the line of statement in module, or None if the line holds others.""" + for line in module.body: + if isinstance(line, cst.SimpleStatementLine) and len(line.body) == 1 and line.body[0] is statement: + return line.trailing_whitespace + return None + + +class _NameCollector(cst.CSTVisitor): + """Collect the names an expression reads: every `Name` except an attribute's or a keyword argument's own name.""" + + def __init__(self) -> None: + """Start with no names collected.""" + super().__init__() + self.names: set[str] = set() + + def visit_Name(self, node: cst.Name) -> None: # noqa: N802 - libcst dispatches on this exact name + """Record the name.""" + self.names.add(node.value) + + def visit_Attribute(self, node: cst.Attribute) -> bool: # noqa: N802 - libcst dispatches on this exact name + """Visit only the attribute's object: in `np.integer`, `integer` is not a name read here.""" + node.value.visit(self) + return False + + def visit_Arg(self, node: cst.Arg) -> bool: # noqa: N802 - libcst dispatches on this exact name + """Visit only the argument's value, not the keyword it is passed as.""" + node.value.visit(self) + return False + + +def names_read_by(node: cst.CSTNode) -> set[str]: + """Return every name `node` reads, built-ins included; attribute names and keyword argument names are not reads.""" + collector = _NameCollector() + node.visit(collector) + return collector.names + + +def assigned_names(statement: cst.Assign | cst.AnnAssign) -> set[str]: + """Return every name an assignment's targets use, including inside tuple or list targets.""" + targets = [target.target for target in statement.targets] if isinstance(statement, cst.Assign) else [statement.target] + collector = _NameCollector() + for target in targets: + target.visit(collector) + return collector.names diff --git a/test/utils/test_cst_utils.py b/test/utils/test_cst_utils.py new file mode 100644 index 00000000..1b50ab22 --- /dev/null +++ b/test/utils/test_cst_utils.py @@ -0,0 +1,96 @@ +# Copyright (C) 2024-2026 TNO +# +# Licensed under the terms specified in the LICENSE file in the root directory +# of this repository, also available at: +# https://github.com/TNO/Renaissance.Py/blob/main/LICENSE + +"""Tests for renaissance.utils.cst_utils.""" + +from typing import cast + +import libcst as cst +import pytest +from hamcrest import assert_that, equal_to, same_instance + +from renaissance.utils.cst_utils import assigned_names, module_statements, names_read_by, trailing_of_sole_statement + + +class TestNamesReadBy: + """Tests for names_read_by.""" + + @pytest.mark.parametrize( + ("expression", "expected"), + [ + pytest.param("x", {"x"}, id="name"), + pytest.param("f(a, key=b)", {"f", "a", "b"}, id="keyword-argument-name-is-not-a-read"), + pytest.param("np.integer", {"np"}, id="attribute-name-is-not-a-read"), + pytest.param("a.b.c(d)", {"a", "d"}, id="attribute-chain-reads-only-its-root"), + pytest.param("int | str", {"int", "str"}, id="built-ins-are-included"), + ], + ) + def test_returns_the_names_an_expression_reads(self, expression: str, expected: set[str]) -> None: + """Verify every name read is returned, but neither an attribute's own name nor a keyword argument's name.""" + assert_that(names_read_by(cst.parse_expression(expression)), equal_to(expected)) + + +class TestAssignedNames: + """Tests for assigned_names.""" + + @pytest.mark.parametrize( + ("statement", "expected"), + [ + pytest.param("x = 1", {"x"}, id="plain"), + pytest.param("x = y = 1", {"x", "y"}, id="chained"), + pytest.param("x: int = 1", {"x"}, id="annotated-leaves-the-annotation-out"), + pytest.param("a, b = 1, 2", {"a", "b"}, id="tuple-target"), + pytest.param("[a, (b, c)] = value", {"a", "b", "c"}, id="nested-targets"), + ], + ) + def test_returns_the_names_the_targets_bind(self, statement: str, expected: set[str]) -> None: + """Verify every name in the targets is returned, including inside tuple or list targets.""" + line = cast("cst.SimpleStatementLine", cst.parse_statement(statement)) + + assert_that(assigned_names(cast("cst.Assign | cst.AnnAssign", line.body[0])), equal_to(expected)) + + +class TestModuleStatements: + """Tests for module_statements.""" + + @pytest.mark.parametrize( + ("source", "expected"), + [ + pytest.param("a = 1\nb = 2\n", ["Assign", "Assign"], id="one-per-line"), + pytest.param("a = 1; b = 2\n", ["Assign", "Assign"], id="one-per-small-statement"), + pytest.param("a = 1\nif a:\n b = 2\n", ["Assign", "If"], id="compound-statement-is-one-item"), + ], + ) + def test_yields_the_top_level_statements_in_order(self, source: str, expected: list[str]) -> None: + """Verify a line of several small statements gives one item each and a compound statement gives one in all.""" + statements = [type(statement).__name__ for statement in module_statements(cst.parse_module(source))] + + assert_that(statements, equal_to(expected)) + + +class TestTrailingOfSoleStatement: + """Tests for trailing_of_sole_statement.""" + + @pytest.mark.parametrize( + "source", + [ + pytest.param("x = 1 # note\n", id="with-comment"), + pytest.param("x = 1\n", id="without-comment"), + ], + ) + def test_returns_the_trailing_whitespace_of_a_sole_statement(self, source: str) -> None: + """Verify the trailing whitespace of the statement's own line is returned.""" + module = cst.parse_module(source) + line = cast("cst.SimpleStatementLine", module.body[0]) + + assert_that(trailing_of_sole_statement(module, line.body[0]), same_instance(line.trailing_whitespace)) + + def test_returns_none_when_the_line_holds_other_statements(self) -> None: + """Verify None is returned for a statement that shares its line with another one.""" + module = cst.parse_module("a = 1; b = 2 # note\n") + line = cast("cst.SimpleStatementLine", module.body[0]) + + assert_that(trailing_of_sole_statement(module, line.body[0]), equal_to(None))