Introduce composition

Author: Julfried

pylint/pyreverse/diagrams.py

@@ -234,15 +234,22 @@ def extract_relationships(self) -> None:
                 except KeyError:
                     continue
 
-            # associations & aggregations links
+            # Composition links
+            for name, values in list(node.compositions_type.items()):
+                for value in values:
+                    self.assign_association_relationship(
+                        value, obj, name, "composition"
+                    )
+
+            # Aggregation links
             for name, values in list(node.aggregations_type.items()):
                 for value in values:
                     self.assign_association_relationship(
                         value, obj, name, "aggregation"
                     )
 
+            # Association links
             associations = node.associations_type.copy()
-
             for name, values in node.locals_type.items():
                 if name not in associations:
                     associations[name] = values

pylint/pyreverse/inspector.py

@@ -122,8 +122,14 @@ def __init__(self, project: Project, tag: bool = False) -> None:
         self.tag = tag
         # visited project
         self.project = project
-        self.associations_handler = AggregationsHandler()
-        self.associations_handler.set_next(OtherAssociationsHandler())
+
+        # Chain: Composition → Aggregation → Association
+        self.associations_handler = CompositionsHandler()
+        aggregation_handler = AggregationsHandler()
+        association_handler = AssociationsHandler()
+
+        self.associations_handler.set_next(aggregation_handler)
+        aggregation_handler.set_next(association_handler)
 
     def visit_project(self, node: Project) -> None:
         """Visit a pyreverse.utils.Project node.
@@ -167,6 +173,7 @@ def visit_classdef(self, node: nodes.ClassDef) -> None:
             specializations.append(node)
             baseobj.specializations = specializations
         # resolve instance attributes
+        node.compositions_type = collections.defaultdict(list)
         node.instance_attrs_type = collections.defaultdict(list)
         node.aggregations_type = collections.defaultdict(list)
         node.associations_type = collections.defaultdict(list)
@@ -327,28 +334,50 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
             self._next_handler.handle(node, parent)
 
 
+class CompositionsHandler(AbstractAssociationHandler):
+    """Handle composition relationships where parent creates child objects."""
+
+    def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
+        if not isinstance(node.parent, (nodes.AnnAssign, nodes.Assign)):
+            super().handle(node, parent)
+            return
+
+        value = node.parent.value
+
+        # Composition: parent creates child (self.x = P())
+        if isinstance(value, nodes.Call):
+            current = set(parent.compositions_type[node.attrname])
+            parent.compositions_type[node.attrname] = list(
+                current | utils.infer_node(node)
+            )
+            return
+
+        # Not a composition, pass to next handler
+        super().handle(node, parent)
+
+
 class AggregationsHandler(AbstractAssociationHandler):
+    """Handle aggregation relationships where parent receives child objects."""
+
     def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
-        # Check if we're not in an assignment context
         if not isinstance(node.parent, (nodes.AnnAssign, nodes.Assign)):
             super().handle(node, parent)
             return
 
         value = node.parent.value
 
-        # Handle direct name assignments
+        # Aggregation: parent receives child (self.x = x)
         if isinstance(value, astroid.node_classes.Name):
             current = set(parent.aggregations_type[node.attrname])
             parent.aggregations_type[node.attrname] = list(
                 current | utils.infer_node(node)
             )
             return
 
-        # Handle comprehensions
+        # Aggregation: comprehensions (self.x = [P() for ...])
         if isinstance(
             value, (nodes.ListComp, nodes.DictComp, nodes.SetComp, nodes.GeneratorExp)
         ):
-            # Determine the type of the element in the comprehension
             if isinstance(value, nodes.DictComp):
                 element_type = safe_infer(value.value)
             else:
@@ -358,12 +387,23 @@ def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
                 parent.aggregations_type[node.attrname] = list(current | {element_type})
                 return
 
-        # Fallback to parent handler
+        # Type annotation only (x: P) defaults to aggregation
+        if isinstance(node.parent, nodes.AnnAssign) and node.parent.value is None:
+            current = set(parent.aggregations_type[node.attrname])
+            parent.aggregations_type[node.attrname] = list(
+                current | utils.infer_node(node)
+            )
+            return
+
+        # Not an aggregation, pass to next handler
         super().handle(node, parent)
 
 
-class OtherAssociationsHandler(AbstractAssociationHandler):
+class AssociationsHandler(AbstractAssociationHandler):
+    """Handle regular association relationships."""
+
     def handle(self, node: nodes.AssignAttr, parent: nodes.ClassDef) -> None:
+        # Everything else is a regular association
         current = set(parent.associations_type[node.attrname])
         parent.associations_type[node.attrname] = list(current | utils.infer_node(node))
 

pylint/pyreverse/printer.py

@@ -22,6 +22,7 @@ class NodeType(Enum):
 
 class EdgeType(Enum):
     INHERITS = "inherits"
+    COMPOSITION = "composition"
     ASSOCIATION = "association"
     AGGREGATION = "aggregation"
     USES = "uses"

pylint/pyreverse/writer.py

@@ -146,6 +146,14 @@ def write_classes(self, diagram: ClassDiagram) -> None:
                 label=rel.name,
                 type_=EdgeType.ASSOCIATION,
             )
+        # generate compositions
+        for rel in diagram.get_relationships("composition"):
+            self.printer.emit_edge(
+                rel.from_object.fig_id,
+                rel.to_object.fig_id,
+                label=rel.name,
+                type_=EdgeType.COMPOSITION,
+            )
         # generate aggregations
         for rel in diagram.get_relationships("aggregation"):
             if rel.to_object.fig_id in associations[rel.from_object.fig_id]: