refactor: add += to Model for Person and SoftwareSystem

Author: yt-ms

src/structurizr/model/model.py

@@ -17,7 +17,7 @@
 
 
 import logging
-from typing import Iterable, List, Optional, ValuesView
+from typing import Iterable, List, Optional, Union, ValuesView
 
 from pydantic import Field
 
@@ -209,6 +209,18 @@ def add_software_system(self, software_system=None, **kwargs) -> SoftwareSystem:
         self.software_systems.add(software_system)
         return software_system
 
+    def __iadd__(self, element: Union[Person, SoftwareSystem]) -> "Model":
+        """Add a new Person or SoftwareSystem to the model."""
+        if isinstance(element, Person):
+            self.add_person(person=element)
+        elif isinstance(element, SoftwareSystem):
+            self.add_software_system(software_system=element)
+        else:
+            raise ValueError(
+                f"Cannot add element with the name {element.name} to Model with += as it is not a Person or a SoftwareSystem."
+            )
+        return self
+
     def add_container(self, container: Container) -> Container:
         """
         Register a newly constructed container with the model.

tests/unit/model/test_model.py

@@ -16,7 +16,7 @@
 
 import pytest
 
-from structurizr.model import Component, Container, Model
+from structurizr.model import Component, Container, Model, Person, SoftwareSystem
 
 
 @pytest.fixture(scope="function")
@@ -76,3 +76,30 @@ def test_model_add_container_must_have_parent(empty_model: Model):
     container = Container(name="c1")
     with pytest.raises(ValueError, match="Container with name .* has no parent"):
         empty_model.add_container(container)
+
+
+def test_model_add_person_with_plusequals(empty_model: Model):
+    """Check that adding a Person to a Model with += works, consistent with += on SoftwareSystem and Container."""
+    bob = Person(name="Bob")
+    empty_model += bob
+    assert bob in empty_model.people
+    assert bob.id != ""
+
+
+def test_model_add_software_system_with_plusequals(empty_model: Model):
+    """Check that adding a SoftwareSystem to a Model with += works, consistent with += on SoftwareSystem and Container."""
+    sys = SoftwareSystem(name="Sys")
+    empty_model += sys
+    assert sys in empty_model.software_systems
+    assert sys.id != ""
+
+
+def test_model_can_only_add_person_or_software_system_with_plusequals(
+    empty_model: Model,
+):
+    """Ensure that passing something other than a Person or SoftwareSystem in to += fails."""
+    c = Container(name="C")
+    with pytest.raises(
+        ValueError, match="Cannot add element with the name .* to Model"
+    ):
+        empty_model += c