refactor: make hydration consistent about adding to the model

Author: yt-ms

src/structurizr/model/component.py

@@ -27,6 +27,7 @@
 
 if TYPE_CHECKING:
     from .container import Container
+    from .model import Model
 
 
 __all__ = ("Component", "ComponentIO")
@@ -108,12 +109,19 @@ def __init__(
         self.tags.add(Tags.COMPONENT)
 
     @classmethod
-    def hydrate(cls, component_io: ComponentIO, container: "Container") -> "Component":
-        """Create and hydrate a new `Component` instance from its IO."""
-        return cls(
+    def hydrate(
+        cls, component_io: ComponentIO, container: "Container", model: "Model"
+    ) -> "Component":
+        """Create and hydrate a new `Component` instance from its IO.
+
+        This will also automatically register with the model.
+        """
+        component = cls(
             **cls.hydrate_arguments(component_io),
             parent=container,
             technology=component_io.technology,
             # TODO: code_elements=map(CodeElement.hydrate, component_io.components),
             size=component_io.size,
         )
+        model += component
+        return component

src/structurizr/model/container.py

@@ -126,7 +126,9 @@ def hydrate(
         model += container
 
         for component_io in container_io.components:
-            component = Component.hydrate(component_io, container=container)
+            component = Component.hydrate(
+                component_io, container=container, model=model
+            )
             container += component
 
         return container

src/structurizr/model/model.py

@@ -149,15 +149,13 @@ def hydrate(cls, model_io: ModelIO) -> "Model":
         )
 
         for person_io in model_io.people:
-            model += Person.hydrate(person_io)
+            Person.hydrate(person_io, model=model)
 
         for software_system_io in model_io.software_systems:
-            model += SoftwareSystem.hydrate(software_system_io, model=model)
+            SoftwareSystem.hydrate(software_system_io, model=model)
 
         for deployment_node_io in model_io.deployment_nodes:
-            DeploymentNode.hydrate(
-                deployment_node_io, model=model
-            )  # Auto-registers with the model
+            DeploymentNode.hydrate(deployment_node_io, model=model)
 
         for element in model.get_elements():
             for relationship in element.relationships:

src/structurizr/model/person.py

@@ -16,7 +16,7 @@
 """Provide a person model."""
 
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from pydantic import Field
 
@@ -26,6 +26,9 @@
 from .tags import Tags
 
 
+if TYPE_CHECKING:
+    from .model import Model
+
 __all__ = ("PersonIO", "Person")
 
 
@@ -60,12 +63,17 @@ def __init__(self, *, location: Location = Location.Unspecified, **kwargs) -> No
         self.tags.add(Tags.PERSON)
 
     @classmethod
-    def hydrate(cls, person_io: PersonIO) -> "Person":
-        """Create a new person and hydrate from its IO."""
-        return cls(
+    def hydrate(cls, person_io: PersonIO, model: "Model") -> "Person":
+        """Create a new person and hydrate from its IO.
+
+        This will also automatically register with the model.
+        """
+        person = cls(
             **cls.hydrate_arguments(person_io),
             location=person_io.location,
         )
+        model += person
+        return person
 
     def interacts_with(
         self, destination: "Person", description: str, **kwargs

src/structurizr/model/software_system.py

@@ -121,11 +121,15 @@ def get_container_with_name(self, name: str) -> Container:
     def hydrate(
         cls, software_system_io: SoftwareSystemIO, model: "Model"
     ) -> "SoftwareSystem":
-        """Create a new SoftwareSystem instance and hydrate it from its IO."""
+        """Create a new SoftwareSystem instance and hydrate it from its IO.
+
+        This will also automatically register with the model.
+        """
         software_system = cls(
             **cls.hydrate_arguments(software_system_io),
             location=software_system_io.location,
         )
+        model += software_system
 
         for container_io in software_system_io.containers:
             Container.hydrate(

tests/unit/model/test_model_item.py

@@ -76,5 +76,5 @@ def test_tag_order_is_preserved_to_and_from_io(empty_model: Model):
     element_io = SoftwareSystemIO.from_orm(element)
     assert element_io.tags == ["Element", "Software System", "tag3", "tag2", "tag1"]
     assert element_io.dict()["tags"] == "Element,Software System,tag3,tag2,tag1"
-    element2 = SoftwareSystem.hydrate(element_io, empty_model)
+    element2 = SoftwareSystem.hydrate(element_io, Model())
     assert list(element2.tags) == ["Element", "Software System", "tag3", "tag2", "tag1"]