Update AST Injection

Author: Behemyth

cppython/plugins/conan/builder.py

@@ -21,7 +21,9 @@ 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]),
+            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:
@@ -46,51 +48,70 @@ def _is_conanfile_class(class_node: cst.ClassDef) -> bool:
 
         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
-        )
+        return any((isinstance(base.value, cst.Name) and base.value.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.
+        """Update or add a 'requires' assignment in a ConanFile class definition."""
+        # Check if 'requires' is already defined
+        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':
+                    continue
+
+                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):
+                break
+            last_attribute = body_statement_line
+
+        # Construct a new statement for the 'requires' attribute
+        new_statement = cst.SimpleStatementLine(
+            body=[self._create_requires_assignment()],
+        )
 
-        Args:
-            updated_node: The class definition to update.
+        # Insert the new statement after the last attribute assignment
+        if last_attribute is not None:
+            new_body = list(updated_node.body.body)
+            index = new_body.index(last_attribute)
+            new_body.insert(index + 1, new_statement)
+        else:
+            new_body = [new_statement] + list(updated_node.body.body)
 
-        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))
-        )
+        return updated_node.with_changes(body=updated_node.body.with_changes(body=new_body))
 
     def _replace_requires(
-        self, updated_node: cst.ClassDef, body_item: cst.SimpleStatementLine, stmt: cst.Assign
+        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.
 
         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.
+            body_statement_line (cst.SimpleStatementLine): The body item containing the assignment.
+            assignment_statement (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)
+        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)
         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]
+                body=[new_assignment if item is body_statement_line else item for item in updated_node.body.body]
             )
         )