Checked some TODOs

functions/lego_classes.py

@@ -2,97 +2,192 @@
 File consists of several classes for the different elements of a device.
 '''
 from __future__ import annotations
+
+from enum import Enum, auto
 import uuid
 from typing import Any, Union, Literal, TypedDict, TypeVar, Type, List, Optional, Dict
+import json
+
 
 # TODO
-# - Parents direkt setzen
-# - Alle gemeinsamen JSON attribute im LegoItem constructor
-# - enforce uuid
 # - Docstrings
 # - Beschreibung von Teilen (-> properties)
 # - Gute String Darstellung
 # - Minimalbeispiel für KPIs
-# - Klassen analog zur deutschen Aufgabenstellung umbenennen (assembly, component) #13 #17
-# - Export als JSON
 # - Export als GraphViz
-# - AggregationsEbene enum (system, assembly, subassembly, component)
 
 
-class LegoItem:
-    def __init__(self, item_number: int, mass: float, delivery_time: int, **kwargs) -> None:
-        # , *args, **kwargs not handling additional/optional specs right now
-        self.id: uuid.UUID = uuid.uuid4()
-        self.properties: dict = kwargs
-        self.item_number: int = item_number
+class ComponentCategory(Enum):
+    BATTERY = auto()
+    MOTOR = auto()
+    FRAME = auto()
+    WHEEL = auto()
+    AXLE = auto()
+    GEAR = auto()
+
+
+class AggregationLayer(Enum):
+    SYSTEM = auto()
+    ASSEMBLY = auto()
+    SUBASSEMBLY = auto()
+    COMPONENT = auto()
+
+
+class LegoComponent:
+    def __init__(self, name: str, category: ComponentCategory, lego_id: str, cost: float, mass: float,
+                 delivery_time: int, layer: AggregationLayer = AggregationLayer.COMPONENT, **properties) -> None:
+        self.uuid: uuid.UUID = uuid.uuid4()
+        self.parent: None | LegoAssembly = None
+        self.name: str = name
+        self.category: ComponentCategory = category
+        self.lego_id: str = lego_id
+        self.cost: float = cost
         self.mass: float = mass
         self.delivery_time: int = delivery_time
-        # TODO: Set parent directly and not via id? This would allow for easier traversal of the tree
-        #   Currently there is no way to search for parts and components in tree by id.
-        self.parent_id = None  # This will be set when added to a component
+        self.layer: AggregationLayer = layer
+        self.properties: dict = properties
+
+    def clone(self) -> LegoComponent:
+        clone = LegoComponent(self.name, self.category, self.lego_id, self.cost, self.mass, self.delivery_time,
+                              self.layer, **self.properties)
+        return clone
+
+    def get_root_assembly(self):
+        if self.parent is None:
+            return None
+        current_assembly = self.parent
+        while current_assembly.parent is not None:
+            current_assembly = current_assembly.parent
+        return current_assembly
+
+    def to_dict(self) -> Dict:
+        ATTRIBUTES = ["uuid", "name", "category", "lego_id", "cost", "mass", "delivery_time", "layer", "properties"]
+        dict_ = {}
+        # store attributes
+        for attr in ATTRIBUTES:
+            dict_[attr] = getattr(self, attr)
 
+        dict_ = {"component": dict_}
+        return dict_
+
+    # TODO good string representation
     def __str__(self):
+        return self.__repr__()
         return (
-            f"Item(id={self.id}, item_number={self.item_number}, "
+            f"Item(id={self.uuid}, item_number={self.lego_id}, "
             f"mass={self.mass}, delivery_time={self.delivery_time}, "
-            f"parent_id={self.parent_id})"
+            f"parent_id={self.parent})"
         )
 
+    # TODO good repr representation
     def __repr__(self):
-        return f"Lego Item [{self.id}]"
+        return f"LegoComponent {self.name} [{self.uuid}]"
 
 
-class LegoComponent:
-    def __init__(self, **kwargs) -> None:
-        self.id: uuid.UUID = uuid.uuid4()
-        self.properties: dict = kwargs
-        self.items: List[LegoItem] = []
+class LegoAssembly:
+    def __init__(self, name: str, layer: AggregationLayer, **properties) -> None:
+        self.uuid: uuid.UUID = uuid.uuid4()
+        self.parent: None | LegoAssembly = None
+        self.name: str = name
+        self.layer: AggregationLayer = layer
+        self.properties: dict = properties
         self.components: List[LegoComponent] = []
-        self.parent_id: None | uuid.UUID = None
-
-    def add_item(self, item: LegoItem) -> None:
-        if not isinstance(item, LegoItem):
-            raise TypeError(f"'item' should be of type LegoPart, got {type(item).__name__} instead.")
-        item.parent_id = self.id
-        self.items.append(item)
+        self.assemblies: List[LegoAssembly] = []
 
     def add_component(self, component: LegoComponent) -> None:
         if not isinstance(component, LegoComponent):
-            raise TypeError(f"'component' should be of type LegoComponent, got {type(component).__name__} instead.")
-        component.parent_id = self.id
+            raise TypeError(f"Argument should be of type {LegoComponent.__name__}, "
+                            f"got {type(component).__name__} instead.")
+        if self.get_root_assembly().contains_uuid(component.uuid):
+            raise AssertionError(f"This assembly or a subassembly already contains the component with ID "
+                                 f"{component.uuid}.")
+        component.parent = self
         self.components.append(component)
 
-    def children(self) -> Dict[str, List[LegoItem] | List[LegoComponent]]:
-        return {'items': self.items, 'components': self.components}
+    def add_assembly(self, assembly: LegoAssembly) -> None:
+        if not isinstance(assembly, LegoAssembly):
+            raise TypeError(f"Argument should be of type {LegoAssembly.__name__}, "
+                            f"got {type(assembly).__name__} instead.")
+        if self.get_root_assembly().contains_uuid(assembly.uuid):
+            raise AssertionError(f"This assembly or a subassembly already contains the assembly with ID "
+                                 f"{assembly.uuid}.")
+        assembly.parent = self
+        self.assemblies.append(assembly)
+
+    def children(self) -> Dict[str, List[LegoComponent] | List[LegoAssembly]]:
+        return {'components': self.components, 'assemblies': self.assemblies}
+
+    def get_component_list(self, max_depth: int = -1) -> List[LegoComponent]:
+        component_list = []
+        component_list.extend(self.components)
+        if max_depth > 0:
+            for assembly in self.assemblies:
+                component_list.extend(assembly.get_component_list(max_depth - 1))
+        return component_list
 
-    def get_item_list(self) -> List[LegoItem]:
-        item_list = []
-        item_list.extend(self.items)
-        for component in self.components:
-            item_list.extend(component.get_item_list())
-        return item_list
+    def get_root_assembly(self) -> LegoAssembly:
+        current_assembly = self
+        while current_assembly.parent is not None:
+            current_assembly = current_assembly.parent
+        return current_assembly
 
+    def contains_uuid(self, uuid_: uuid.UUID):
+        # check component ids
+        component_ids = list(map(lambda c: c.uuid, self.components))
+        if uuid_ in component_ids:
+            return True
+        # check assembly ids
+        assembly_ids = list(map(lambda a: a.uuid, self.assemblies))
+        if uuid_ in assembly_ids:
+            return True
+        # recursively check assemblies
+        for assembly in self.assemblies:
+            if assembly.contains_uuid(uuid_):
+                return True
+        return False
+
+    def to_dict(self) -> Dict:
+        ATTRIBUTES = ["uuid", "name", "layer", "properties"]
+        dict_ = {}
+        # store attributes
+        for attr in ATTRIBUTES:
+            dict_[attr] = getattr(self, attr)
+        # store components
+        dict_["components"] = [component.to_dict() for component in self.components]
+        dict_["assemblies"] = [assembly.to_dict() for assembly in self.assemblies]
+
+        return {"assembly": dict_}
+
+    # TODO find good string representation
     def __repr__(self):
-        return f"Lego Component [{self.id}]"
+        return f"LegoAssembly {self.name} [{self.uuid}]"
 
-# TODO: Adjust default output when printing an item or component
 
-def print_component_tree(root, level=0, is_last=False):
+def print_assembly_tree(root, level=0, is_last=False):
     # print component
-    component_padding = ""
+    assembly_padding = ""
     if level > 0:
-        component_padding += "│   " * (level - 1)
+        assembly_padding += "│   " * (level - 1)
         if is_last:
-            component_padding += "└── "
+            assembly_padding += "└── "
         else:
-            component_padding += "├── "
-    print(f"{component_padding}{root}")
+            assembly_padding += "├── "
+    print(f"{assembly_padding}{root}")
     # recursively print child components
-    for i, component in enumerate(root.components):
-        is_last_ = i == len(root.components) - 1 and len(root.items) == 0
-        print_component_tree(component, level + 1, is_last_)
+    for i, assembly in enumerate(root.assemblies):
+        is_last_ = i == len(root.assemblies) - 1 and len(root.components) == 0
+        print_assembly_tree(assembly, level + 1, is_last_)
     # print items
-    for i, item in enumerate(root.items):
-        item_padding = "│   " * level
-        item_padding += "├── " if i < len(root.items) - 1 else "└── "
-        print(f"{item_padding}{item}")
+    for i, item in enumerate(root.components):
+        component_padding = "│   " * level if not is_last else "    "
+        component_padding += "├── " if i < len(root.components) - 1 else "└── "
+        print(f"{component_padding}{item}")
+
+
+class KPIEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, uuid.UUID):
+            return "kpi-" + str(o)
+        if isinstance(o, (ComponentCategory, AggregationLayer)):
+            return "kpi-" + o.name
+        return super().default(o)