Add Test + Fix Ast Rewrite

Author: Behemyth

cppython/plugins/conan/builder.py

@@ -20,10 +20,10 @@ def __init__(self, dependencies: list[ConanDependency]) -> None:
     def _create_requires_assignment(self) -> cst.Assign:
         """Create a `requires` assignment statement."""
         return cst.Assign(
-            targets=[cst.AssignTarget(cst.Name('requires'))],
-            value=cst.List([
-                cst.Element(cst.SimpleString(f'"{dependency.requires()}"')) for dependency in self.dependencies
-            ]),
+            targets=[cst.AssignTarget(cst.Name(value='requires'))],
+            value=cst.List(
+                [cst.Element(cst.SimpleString(f'"{dependency.requires()}"')) for dependency in self.dependencies]
+            ),
         )
 
     def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.BaseStatement:
@@ -56,24 +56,23 @@ def _update_requires(self, updated_node: cst.ClassDef) -> cst.ClassDef:
         for body_statement_line in updated_node.body.body:
             if not isinstance(body_statement_line, cst.SimpleStatementLine):
                 continue
-
-            assignment_statement = body_statement_line.body[0]
-            if not isinstance(assignment_statement, cst.Assign):
-                continue
-
-            for target in assignment_statement.targets:
-                if not isinstance(target.target, cst.Name) or target.target.value != 'requires':
+            for assignment_statement in body_statement_line.body:
+                if not isinstance(assignment_statement, cst.Assign):
                     continue
-
-                return self._replace_requires(updated_node, body_statement_line, assignment_statement)
+                for target in assignment_statement.targets:
+                    if not isinstance(target.target, cst.Name) or target.target.value != 'requires':
+                        continue
+                    # Replace only the assignment within the SimpleStatementLine
+                    return self._replace_requires(updated_node, body_statement_line, assignment_statement)
 
         # Find the last attribute assignment before methods
         last_attribute = None
         for body_statement_line in updated_node.body.body:
             if not isinstance(body_statement_line, cst.SimpleStatementLine):
                 break
-            assignment_statement = body_statement_line.body[0]
-            if not isinstance(assignment_statement, cst.Assign):
+            if not body_statement_line.body:
+                break
+            if not isinstance(body_statement_line.body[0], cst.Assign):
                 break
             last_attribute = body_statement_line
 
@@ -89,29 +88,27 @@ def _update_requires(self, updated_node: cst.ClassDef) -> cst.ClassDef:
             new_body.insert(index + 1, new_statement)
         else:
             new_body = [new_statement] + list(updated_node.body.body)
-
         return updated_node.with_changes(body=updated_node.body.with_changes(body=new_body))
 
     def _replace_requires(
         self, updated_node: cst.ClassDef, body_statement_line: cst.SimpleStatementLine, assignment_statement: cst.Assign
     ) -> cst.ClassDef:
-        """Replace the existing 'requires' assignment with a new one.
+        """Replace the existing 'requires' assignment with a new one, preserving other statements on the same line."""
+        new_value = cst.List(
+            [cst.Element(cst.SimpleString(f'"{dependency.requires()}"')) for dependency in self.dependencies]
+        )
+        new_assignment = assignment_statement.with_changes(value=new_value)
 
-        Args:
-            updated_node (cst.ClassDef): The class definition to update.
-            body_statement_line (cst.SimpleStatementLine): The body item containing the assignment.
-            assignment_statement (cst.Assign): The existing assignment statement.
+        # Replace only the relevant assignment in the SimpleStatementLine
+        new_body = [
+            new_assignment if statement is assignment_statement else statement for statement in body_statement_line.body
+        ]
+        new_statement_line = body_statement_line.with_changes(body=new_body)
 
-        Returns:
-            cst.ClassDef: The updated class definition.
-        """
-        new_value = cst.List([
-            cst.Element(cst.SimpleString(f'"{dependency.requires()}"')) for dependency in self.dependencies
-        ])
-        new_assignment = assignment_statement.with_changes(value=new_value)
+        # Replace the statement line in the class body
         return updated_node.with_changes(
             body=updated_node.body.with_changes(
-                body=[new_assignment if item is body_statement_line else item for item in updated_node.body.body]
+                body=[new_statement_line if item is body_statement_line else item for item in updated_node.body.body]
             )
         )
 

tests/unit/plugins/conan/test_ast.py

@@ -0,0 +1,80 @@
+"""Tests for the AST transformer that modifies ConanFile classes."""
+
+import ast
+from textwrap import dedent
+
+import libcst as cst
+
+from cppython.plugins.conan.builder import RequiresTransformer
+from cppython.plugins.conan.schema import ConanDependency
+
+
+class TestTransformer:
+    """Unit tests for the RequiresTransformer."""
+
+    class MockDependency(ConanDependency):
+        """A dummy dependency class for testing."""
+
+        @staticmethod
+        def requires() -> str:
+            """Return a dummy requires string."""
+            return 'test/1.2.3'
+
+    @staticmethod
+    def test_add_requires_when_missing() -> None:
+        """Test that the transformer adds requires when missing."""
+        dependency = TestTransformer.MockDependency(name='test')
+
+        code = """
+        class MyFile(ConanFile):
+            name = "test"
+            version = "1.0"
+        """
+
+        module = cst.parse_module(dedent(code))
+        transformer = RequiresTransformer([dependency])
+        modified = module.visit(transformer)
+        assert 'requires = ["test/1.2.3"]' in modified.code
+
+        # Verify the resulting code is valid Python syntax
+        ast.parse(modified.code)
+
+    @staticmethod
+    def test_replace_existing_requires() -> None:
+        """Test that the transformer replaces existing requires."""
+        dependency = TestTransformer.MockDependency(name='test')
+
+        code = """
+        class MyFile(ConanFile):
+            name = "test"
+            requires = ["old/0.1"]
+            version = "1.0"
+        """
+
+        module = cst.parse_module(dedent(code))
+        transformer = RequiresTransformer([dependency])
+        modified = module.visit(transformer)
+        assert 'requires = ["test/1.2.3"]' in modified.code
+        assert 'old/0.1' not in modified.code
+
+        # Verify the resulting code is valid Python syntax
+        ast.parse(modified.code)
+
+    @staticmethod
+    def test_no_conanfile_class() -> None:
+        """Test that the transformer does not modify non-ConanFile classes."""
+        dependency = TestTransformer.MockDependency(name='test')
+
+        code = """
+        class NotConan:
+            pass
+        """
+
+        module = cst.parse_module(dedent(code))
+        transformer = RequiresTransformer([dependency])
+        modified = module.visit(transformer)
+        # Should not add requires to non-ConanFile classes
+        assert 'requires' not in modified.code
+
+        # Verify the resulting code is valid Python syntax
+        ast.parse(modified.code)