diff --git a/functions/lego_classes.py b/functions/lego_classes.py index 7e05af77dd38cb6f0f21b56b9e5f56dd380ed231..92ed684dc55261b904d5f255340974d880397934 100644 --- a/functions/lego_classes.py +++ b/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