diff --git a/functions/classes.py b/functions/classes.py index 6e315c7e37adb4a0a55e6fcc1ab660589d2be172..9c0d9fe17a72461ab32c36697fd890dc188903ec 100644 --- a/functions/classes.py +++ b/functions/classes.py @@ -65,24 +65,13 @@ class LegoComponent: 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_ + dict_ = { + "uuid": self.uuid, + "label": self.label, + "properties": self.properties, + "layer": self.layer, + } + return {"component": dict_} # TODO good string representation def __str__(self): @@ -169,15 +158,15 @@ class LegoAssembly: return False def to_dict(self) -> Dict: - ATTRIBUTES = ["uuid", "name", "layer", "properties"] - dict_ = {} - # store attributes - for attr in ATTRIBUTES: - dict_[attr] = getattr(self, attr) + dict_ = { + "uuid": self.uuid, + "label": self.label, + "properties": self.properties, + "layer": self.layer, + } # 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