Update libcst utils used in recipes

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

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))