diff --git a/app/cli/db.py b/app/cli/db.py index 72ea5eb..c0f2661 100644 --- a/app/cli/db.py +++ b/app/cli/db.py @@ -45,18 +45,16 @@ models = [ class ExportEncoder(json.JSONEncoder): """Encoder to serialise all types used in the database.""" - def default(self, obj: Any) -> Any: - if isinstance(obj, AlarmState): - return obj.name - elif isinstance(obj, AutomationState): - return obj.name - elif isinstance(obj, bytes): - return base64.encodebytes(obj).decode('utf-8') - elif isinstance(obj, (datetime.datetime, datetime.date, datetime.time)): - return obj.isoformat() - elif isinstance(obj, datetime.timedelta): - return (datetime.datetime.min + obj).time().isoformat() - return super().default(obj) + def default(self, o: Any) -> Any: + if isinstance(o, AlarmState): + return o.name + if isinstance(o, AutomationState): + return o.name + if isinstance(o, bytes): + return base64.encodebytes(o).decode('utf-8') + if isinstance(o, (datetime.datetime, datetime.date, datetime.time)): + return o.isoformat() + return super().default(o) def model_to_dict(model: db.Model) -> Dict[str, Any]: # type: ignore[name-defined] @@ -80,10 +78,9 @@ decoder = { "AlarmState": lambda x: AlarmState.__getattribute__(AlarmState, x), "AutomationState": lambda x: AutomationState.__getattribute__(AutomationState, x), "bytes": lambda x: base64.decodebytes(x.encode('utf-8')), - "datetime": lambda x: datetime.datetime.fromisoformat(x), - "int": lambda x: int(x), + "datetime": datetime.datetime.fromisoformat, + "int": int, "str": lambda x: x, - # TODO: timedelta (not currently used but could be in the future) } @@ -92,7 +89,7 @@ def db_import_model(model: db.Model, data: List[Dict[str, Any]]) -> None: # typ new = model() for col in row: type_name, col_name = col.split("_", 1) - new.__setattr__(col_name, decoder.get(type_name, lambda x: x)(row[col])) # type: ignore[no-untyped-call] + setattr(new, col_name, decoder.get(type_name, lambda x: x)(row[col])) # type: ignore[no-untyped-call] db.session.add(new)