Refactor LegoComponent

functions/lego_classes.py

@@ -28,49 +28,38 @@ class LegoItem:
     def __repr__(self):
         return f"Lego Item [{self.id}]"
 
-class LegoComponent:
-    def __init__(self, items=None | LegoItem | list[LegoItem],
-                 components=None) -> None:  # TODO: type hinting
-        self.id = uuid.uuid4()
-        if (isinstance(items, list) and items
-           and isinstance(items[0], LegoItem)):
-            self.items = items
-        elif isinstance(items, LegoItem):
-            self.items = [items]
-        else:
-            self.items = []  # must be a list for now,
-        if components is None:
-            self.components = []
-        elif isinstance(components, list):
-            self.components = components
-        else:
-            self.components = [components]  # Not checking correct type
-        self.parent_id = None  # This will be set when added to a component
 
-    def add_item(self, item: LegoItem | List[LegoItem]) -> None:
-        if isinstance(item, list):
-            self.items.extend(item)
-        else:
-            self.items.append(item)
+class LegoComponent:
+    def __init__(self, **kwargs) -> None:
+        self.id: uuid.UUID = uuid.uuid4()
+        self.properties: dict = kwargs
+        self.items: List[LegoItem] = []
+        self.components: List[LegoComponent] = []
+        self.parent_id: None | uuid.UUID = None
 
-    def add_component(self, component: LegoComponent | List[LegoComponent]
-                      ) -> None:
-        if isinstance(component, list):
-            self.components.extend(component)
-        else:
-            self.components.append(component)
+    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)
 
-    def add(self, element) -> None:
-        if isinstance(element, LegoItem):
-            self.items.append(element)
-        elif isinstance(element, LegoComponent):
-            self.components.append(element)
-        else:
-            raise TypeError(f"added elements should be of type LegoItem or "
-                            f"LegoComponent not {type(element)}")
+    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
+        self.components.append(component)
 
     def children(self) -> Dict[str, List[LegoItem] | List[LegoComponent]]:
         return {'items': self.items, 'components': self.components}
 
+    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 __repr__(self):
+        return f"Lego Component [{self.id}]"
 
-# TODO: Adjust default output when printing an item or component
\ No newline at end of file
+# TODO: Adjust default output when printing an item or component