AST Modification

Author: Behemyth

cppython/plugins/conan/builder.py

@@ -4,11 +4,97 @@
 from string import Template
 from textwrap import dedent
 
+import libcst as cst
 from pydantic import DirectoryPath
 
 from cppython.plugins.conan.schema import ConanDependency
 
 
+class RequiresTransformer(cst.CSTTransformer):
+    """Transformer to add or update the `requires` attribute in a ConanFile class."""
+
+    def __init__(self, dependencies: list[ConanDependency]) -> None:
+        """Initialize the transformer with a list of dependencies."""
+        self.dependencies = dependencies
+
+    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'"{dep.requires()}"')) for dep in self.dependencies]),
+        )
+
+    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.BaseStatement:
+        """Modify the class definition to include or update 'requires'.
+
+        Args:
+            original_node: The original class definition.
+            updated_node: The updated class definition.
+
+        Returns: The modified class definition.
+        """
+        if self._is_conanfile_class(original_node):
+            updated_node = self._update_requires(updated_node)
+        return updated_node
+
+    @staticmethod
+    def _is_conanfile_class(class_node: cst.ClassDef) -> bool:
+        """Check if the class inherits from ConanFile.
+
+        Args:
+            class_node: The class definition to check.
+
+        Returns: True if the class inherits from ConanFile, False otherwise.
+        """
+        return any(
+            (isinstance(base, cst.Name) and base.value == 'ConanFile')
+            or (isinstance(base, cst.Attribute) and base.attr.value == 'ConanFile')
+            for base in class_node.bases
+        )
+
+    def _update_requires(self, updated_node: cst.ClassDef) -> cst.ClassDef:
+        """Update or add the 'requires' attribute in the class definition.
+
+        Args:
+            updated_node: The class definition to update.
+
+        Returns: The updated class definition.
+        """
+        for body_item in updated_node.body.body:
+            if isinstance(body_item, cst.SimpleStatementLine):
+                stmt = body_item.body[0]
+                if isinstance(stmt, cst.Assign):
+                    for target in stmt.targets:
+                        if isinstance(target.target, cst.Name) and target.target.value == 'requires':
+                            return self._replace_requires(updated_node, body_item, stmt)
+
+        new_stmt = self._create_requires_assignment()
+        return updated_node.with_changes(
+            body=updated_node.body.with_changes(body=[new_stmt] + list(updated_node.body.body))
+        )
+
+    def _replace_requires(
+        self, updated_node: cst.ClassDef, body_item: cst.SimpleStatementLine, stmt: cst.Assign
+    ) -> cst.ClassDef:
+        """Replace the existing 'requires' assignment with a new one.
+
+        Args:
+            updated_node (cst.ClassDef): The class definition to update.
+            body_item (cst.SimpleStatementLine): The body item containing the assignment.
+            stmt (cst.Assign): The existing assignment statement.
+
+        Returns:
+            cst.ClassDef: The updated class definition.
+        """
+        new_value = cst.List([cst.Element(cst.SimpleString(f'"{dep.requires()}"')) for dep in self.dependencies])
+        new_stmt = stmt.with_changes(value=new_value)
+        return updated_node.with_changes(
+            body=updated_node.body.with_changes(
+                body=[new_stmt if item is body_item else item for item in updated_node.body.body]
+            )
+        )
+
+
 class Builder:
     """Aids in building the information needed for the Conan plugin"""
 
@@ -57,12 +143,14 @@ def generate_conanfile(self, directory: DirectoryPath, dependencies: list[ConanD
         """Generate a conanfile.py file for the project."""
         conan_file = directory / self._filename
 
-        # If the file exists then we need to inject our information into it
         if conan_file.exists():
-            raise NotImplementedError(
-                'Updating existing conanfile.py is not yet supported. Please remove the file and try again.'
-            )
+            source_code = conan_file.read_text(encoding='utf-8')
+
+            module = cst.parse_module(source_code)
+            transformer = RequiresTransformer(dependencies)
+            modified = module.visit(transformer)
 
+            conan_file.write_text(modified.code, encoding='utf-8')
         else:
             directory.mkdir(parents=True, exist_ok=True)
             self._create_conanfile(conan_file, dependencies)

tests/integration/examples/test_conan_cmake.py

@@ -36,3 +36,23 @@ def test_simple(example_runner: CliRunner) -> None:
 
         # Verify that the build directory contains the expected files
         assert (Path('build') / 'CMakeCache.txt').exists(), 'build/CMakeCache.txt not found'
+
+    @staticmethod
+    def test_inject(example_runner: CliRunner) -> None:
+        """Inject"""
+        result = example_runner.invoke(
+            app,
+            [
+                'install',
+            ],
+        )
+
+        assert result.exit_code == 0, result.output
+
+        # Run the CMake configuration command
+        cmake_result = subprocess.run(['cmake', '--preset=default'], capture_output=True, text=True, check=False)
+
+        assert cmake_result.returncode == 0, f'CMake configuration failed: {cmake_result.stderr}'
+
+        # Verify that the build directory contains the expected files
+        assert (Path('build') / 'CMakeCache.txt').exists(), 'build/CMakeCache.txt not found'