import inspect from tinydb import TinyDB, table from tinydb.table import Document from grung.types import Record class RecordTable(table.Table): """ Wrapper around tinydb Tables that handles Records instead of dicts. """ def __init__(self, storage, name, document_class: Document = Record, **kwargs): self.document_class = document_class super().__init__(storage, name, **kwargs) def insert(self, document): self._satisfy_constraints(document) if document.doc_id: last_insert_id = super().upsert(document)[0] else: last_insert_id = super().insert(dict(document)) return self.get(doc_id=last_insert_id) def remove(self, document): if document.doc_id: super().remove(doc_ids=[document.doc_id]) def _satisfy_constraints(self, document): # check for uniqueness, etc. pass class GrungDB(TinyDB): """ A TinyDB database instance that uses RecordTable instances for each table and Record instances for each document in the table. """ default_table_name = "Record" _tables = {} def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.create_table(Record) def table(self, name: str) -> RecordTable: if name not in self._tables: raise RuntimeError(f"No such table: {name}") return self._tables[name] def create_table(self, table_class): name = table_class.__name__ if name not in self._tables: self._tables[name] = RecordTable(self.storage, name, document_class=table_class) return self.table(name) def save(self, record): """ Create or update a record in its table. """ return self.table(record._metadata.table).insert(record) def delete(self, record): return self.table(record._metadata.table).remove(record) def __getattr__(self, attr_name): """ Make tables attributes of the instance. """ if attr_name in self._tables: return self.table(attr_name) return super().__getattr__(attr_name) @classmethod def with_schema(cls, schema_module, *args, **kwargs): db = GrungDB(*args, **kwargs) for name, obj in inspect.getmembers(schema_module): if type(obj) == type and issubclass(obj, Record): db.create_table(obj) return db