add type annotations

pystructurizr/dsl.py

@@ -1,11 +1,14 @@
 import re
 import keyword
+from enum import Enum
+from typing import List, Optional
+
 
 class Identifier:
     counter = {}
 
     @staticmethod
-    def make_identifier(name):
+    def make_identifier(name: str) -> str:
         # Replace invalid characters
         identifier = re.sub('[^0-9a-zA-Z_]', '_', name.lower())
 
@@ -35,77 +38,47 @@ def __init__(self):
         self.level = 0
         self.lines = []
 
-    def add(self, str):
-        self.lines.append(f'{"  " * self.level}{str}')
+    def add(self, s: str) -> None:
+        self.lines.append(f'{"  " * self.level}{s}')
 
-    def indent(self):
+    def indent(self) -> None:
         self.level += 1
 
     def outdent(self):
         self.level -= 1
         if self.level < 0:
             self.level = 0
 
-    def result(self):
+    def result(self) -> str:
         return "\n".join(self.lines)
 
 
 class Element:
-    def __init__(self, name, description=None, technology=None, tags=None):
+    def __init__(self, name: str, description: Optional[str]=None, technology: Optional[str]=None, tags: Optional[List[str]]=None):
         self.name = name
         self.description = description
         self.technology = technology
         self.tags = tags if tags else []
         self.relationships = []
         self.instname = Identifier.make_identifier(name)
 
-    def uses(self, destination, description=None, technology=None):
+    def uses(self, destination: str, description: Optional[str]=None, technology: Optional[str]=None) -> 'Relationship':
         relationship = Relationship(self, destination, description, technology)
         self.relationships.append(relationship)
         return relationship
 
-    def dump(self, dumper):
+    def dump(self, dumper: Dumper) -> None:
         raise NotImplementedError("This method must be implemented in a subclass.")
 
-    def dump_relationships(self, dumper):
+    def dump_relationships(self, dumper: Dumper) -> None:
         raise NotImplementedError("This method must be implemented in a subclass.")
 
 
-class Model(Element):
-    def __init__(self, name):
-        super().__init__(name)
-        self.elements = []
-
-    def Person(self, *args, **kwargs):
-        if args and isinstance(args[0], Person):
-            person = args[0]
-        else:
-            person = Person(*args, **kwargs)
-        self.elements.append(person)
-        return person
-
-    def SoftwareSystem(self, *args, **kwargs):
-        if args and isinstance(args[0], SoftwareSystem):
-            system = args[0]
-        else:
-            system = SoftwareSystem(*args, **kwargs)
-        self.elements.append(system)
-        return system
-
-    def dump(self, dumper):
-        for element in self.elements:
-            element.dump(dumper)
-
-    def dump_relationships(self, dumper):
-        for element in self.elements:
-            element.dump_relationships(dumper)
-
-
 class Person(Element):
-    def __init__(self, name, description=None, technology=None, tags=None):
+    def __init__(self, name: str, description: Optional[str]=None, technology: Optional[str]=None, tags: Optional[List[str]]=None):
         super().__init__(name, description, technology, tags)
 
-    def dump(self, dumper):
+    def dump(self, dumper: Dumper) -> None:
         dumper.add(f'{self.instname} = Person "{self.name}" "{self.description}" {{')
         dumper.indent()
         if self.technology:
@@ -115,57 +88,44 @@ def dump(self, dumper):
         dumper.outdent()
         dumper.add(f'}}')
 
-    def dump_relationships(self, dumper):
+    def dump_relationships(self, dumper: Dumper) -> None:
         for rel in self.relationships:
             rel.dump(dumper)
 
 
-class SoftwareSystem(Element):
-    def __init__(self, name, description=None, technology=None, tags=None):
+class Component(Element):
+    def __init__(self, name: str, description: Optional[str]=None, technology: Optional[str]=None, tags: Optional[List[str]]=None):
         super().__init__(name, description, technology, tags)
-        self.containers = []
 
-    def Container(self, *args, **kwargs):
-        if args and isinstance(args[0], Container):
-            container = args[0]
-        else:
-            container = Container(*args, **kwargs)
-        self.containers.append(container)
-        return container
-
-    def dump(self, dumper):
-        dumper.add(f'{self.instname} = SoftwareSystem "{self.name}" "{self.description}" {{')
+    def dump(self, dumper: Dumper) -> None:
+        dumper.add(f'{self.instname} = Component "{self.name}" "{self.description}" {{')
         dumper.indent()
         if self.technology:
             dumper.add(f'technology "{self.technology}"')
         if self.tags:
             dumper.add(f'tags "{", ".join(self.tags)}"')
-        for container in self.containers:
-            container.dump(dumper)
         dumper.outdent()
         dumper.add(f'}}')
 
-    def dump_relationships(self, dumper):
+    def dump_relationships(self, dumper: Dumper) -> None:
         for rel in self.relationships:
             rel.dump(dumper)
-        for container in self.containers:
-            container.dump_relationships(dumper)
 
 
 class Container(Element):
-    def __init__(self, name, description=None, technology=None, tags=None):
+    def __init__(self, name: str, description: Optional[str]=None, technology: Optional[str]=None, tags: Optional[List[str]]=None):
         super().__init__(name, description, technology, tags)
         self.components = []
 
-    def Component(self, *args, **kwargs):
+    def Component(self, *args, **kwargs) -> Component:
         if args and isinstance(args[0], Component):
             component = args[0]
         else:
             component = Component(*args, **kwargs)
         self.components.append(component)
         return component
 
-    def dump(self, dumper):
+    def dump(self, dumper: Dumper) -> None:
         dumper.add(f'{self.instname} = Container "{self.name}" "{self.description}" {{')
         dumper.indent()
         if self.technology:
@@ -177,60 +137,109 @@ def dump(self, dumper):
         dumper.outdent()
         dumper.add(f'}}')
 
-    def dump_relationships(self, dumper):
+    def dump_relationships(self, dumper: Dumper) -> None:
         for rel in self.relationships:
             rel.dump(dumper)
         for component in self.components:
             component.dump_relationships(dumper)
 
 
-class Component(Element):
-    def __init__(self, name, description=None, technology=None, tags=None):
+class SoftwareSystem(Element):
+    def __init__(self, name: str, description: Optional[str]=None, technology: Optional[str]=None, tags: Optional[List[str]]=None):
         super().__init__(name, description, technology, tags)
+        self.containers = []
 
-    def dump(self, dumper):
-        dumper.add(f'{self.instname} = Component "{self.name}" "{self.description}" {{')
+    def Container(self, *args, **kwargs) -> Container:
+        if args and isinstance(args[0], Container):
+            container = args[0]
+        else:
+            container = Container(*args, **kwargs)
+        self.containers.append(container)
+        return container
+
+    def dump(self, dumper: Dumper) -> None:
+        dumper.add(f'{self.instname} = SoftwareSystem "{self.name}" "{self.description}" {{')
         dumper.indent()
         if self.technology:
             dumper.add(f'technology "{self.technology}"')
         if self.tags:
             dumper.add(f'tags "{", ".join(self.tags)}"')
+        for container in self.containers:
+            container.dump(dumper)
         dumper.outdent()
         dumper.add(f'}}')
 
-    def dump_relationships(self, dumper):
+    def dump_relationships(self, dumper: Dumper) -> None:
         for rel in self.relationships:
             rel.dump(dumper)
+        for container in self.containers:
+            container.dump_relationships(dumper)
+
+
+class Model(Element):
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.elements = []
+
+    def Person(self, *args, **kwargs) -> Person:
+        if args and isinstance(args[0], Person):
+            person = args[0]
+        else:
+            person = Person(*args, **kwargs)
+        self.elements.append(person)
+        return person
+
+    def SoftwareSystem(self, *args, **kwargs) -> SoftwareSystem:
+        if args and isinstance(args[0], SoftwareSystem):
+            system = args[0]
+        else:
+            system = SoftwareSystem(*args, **kwargs)
+        self.elements.append(system)
+        return system
+
+    def dump(self, dumper: Dumper) -> None:
+        for element in self.elements:
+            element.dump(dumper)
+
+    def dump_relationships(self, dumper: Dumper) -> None:
+        for element in self.elements:
+            element.dump_relationships(dumper)
 
 
 class Relationship:
-    def __init__(self, source, destination, description=None, technology=None):
+    def __init__(self, source: Element, destination: Element, description: Optional[str]=None, technology: Optional[str]=None):
         self.source = source
         self.destination = destination
         self.description = description
         self.technology = technology
 
-    def dump(self, dumper):
+    def dump(self, dumper: Dumper) -> None:
         dumper.add(f'{self.source.instname} -> {self.destination.instname} "{self.description}" "{self.technology}"')
 
 
 class View:
-    def __init__(self, viewkind, element, name, description=None):
+    class Kind(Enum):
+        SYSTEM_LANDSCAPE = "systemLandscape"
+        SYSTEM_CONTEXT = "systemContext"
+        CONTAINER = "container"
+        COMPONENT = "component"
+
+    def __init__(self, viewkind: Kind, element: Element, name: str, description: Optional[str]=None):
         self.viewkind = viewkind
         self.element = element
         self.name = name
         self.description = description
         self.includes = []
         self.excludes = []
 
-    def include(self, element):
+    def include(self, element: Element) -> None:
         self.includes.append(element)
 
-    def exclude(self, element):
+    def exclude(self, element: Element) -> None:
         self.excludes.append(element)
 
-    def dump(self, dumper):
-        dumper.add(f'{self.viewkind} {self.element.instname if self.element else ""} {{')
+    def dump(self, dumper: Dumper) -> None:
+        dumper.add(f'{self.viewkind.value} {self.element.instname if self.element else ""} {{')
         dumper.indent()
         if self.description:
             dumper.add(f'description "{self.description}"')
@@ -245,10 +254,10 @@ def dump(self, dumper):
 
 
 class Style:
-    def __init__(self, map):
+    def __init__(self, map: dict[str, str]):
         self.map = map
 
-    def dump(self, dumper):
+    def dump(self, dumper: Dumper) -> None:
         dumper.add(f'element "{self.map["tag"]}" {{')
         dumper.indent()
         for k, v in self.map.items():
@@ -265,7 +274,7 @@ def __init__(self):
         self.views = []
         self.styles = []
 
-    def dump(self, dumper = Dumper()):
+    def dump(self, dumper: Dumper = Dumper()) -> None:
         dumper.add(f'workspace {{')
         dumper.indent()
 
@@ -295,33 +304,33 @@ def dump(self, dumper = Dumper()):
         dumper.add(f'}}')
         return dumper.result()
 
-    def Model(self, model=None, name=None):
+    def Model(self, model: Optional[Model]=None, name: Optional[str]=None):
         if model is None:
             model = Model(name)
         self.models.append(model)
         return model
 
-    def SystemLandscapeView(self, name, description):
-        view = View("systemLandscape", None, name, description)
+    def SystemLandscapeView(self, name: str, description: str):
+        view = View(View.Kind.SYSTEM_LANDSCAPE, None, name, description)
         self.views.append(view)
         return view
 
-    def SystemContextView(self, element, name, description):
-        view = View("systemContext", element, name, description)
+    def SystemContextView(self, element: Element, name: str, description: str):
+        view = View(View.Kind.SYSTEM_CONTEXT, element, name, description)
         self.views.append(view)
         return view
 
-    def ContainerView(self, element, name, description):
-        view = View("container", element, name, description)
+    def ContainerView(self, element: Element, name: str, description: str):
+        view = View(View.Kind.CONTAINER, element, name, description)
         self.views.append(view)
         return view
 
-    def ComponentView(self, element, name, description):
-        view = View("component", element, name, description)
+    def ComponentView(self, element: Element, name: str, description: str):
+        view = View(View.Kind.COMPONENT, element, name, description)
         self.views.append(view)
         return view
 
-    def Styles(self, *styles):
+    def Styles(self, *styles: dict[str, str]) -> None:
         for style in styles:
             self.styles.append(Style(style))