|
| 1 | +from __future__ import annotations |
| 2 | +from abc import ABC |
| 3 | +from typing import get_type_hints, get_origin, get_args, TypeVar, Tuple, Type, Generic |
| 4 | +from types import UnionType, NoneType |
| 5 | +from . import io_safe |
| 6 | +from .sessions import SessionFileFull, SessionFileKey, SessionFileWhere, SessionDirFull, SessionDirWhere |
| 7 | + |
| 8 | + |
| 9 | + |
| 10 | +T = TypeVar("T") |
| 11 | +T2 = TypeVar("T2") |
| 12 | + |
| 13 | + |
| 14 | +def get_type_hints_excluding_internals(cls): |
| 15 | + """ |
| 16 | + Get type hints of the class, excluding double dunder variables. |
| 17 | + """ |
| 18 | + for var_name, var_type in get_type_hints(cls).items(): |
| 19 | + if var_name.startswith("__") and var_name.endswith("__"): |
| 20 | + continue |
| 21 | + yield var_name, var_type |
| 22 | + |
| 23 | + |
| 24 | + |
| 25 | +def fill_object_from_dict_using_type_hints(obj, cls, data: dict): |
| 26 | + """ |
| 27 | + Attributes of obj are set using the data dict. |
| 28 | + The type hints of the class cls are used to determine which attributes to set. |
| 29 | + """ |
| 30 | + for var_name, var_type in get_type_hints_excluding_internals(cls): |
| 31 | + var_type_args = get_args(var_type) |
| 32 | + var_type_origin = get_origin(var_type) |
| 33 | + # Check if variable is nullable (e.g. email: str | None) |
| 34 | + # When it is not nullable but not in the data, raise an error |
| 35 | + if var_name not in data: |
| 36 | + nullable = var_type_origin is UnionType and NoneType in var_type_args |
| 37 | + if not nullable: |
| 38 | + raise KeyError(f"Missing variable '{var_name}' in {cls.__name__}.") |
| 39 | + # When it is a list, fill the list with the items |
| 40 | + if var_type_origin is list and len(var_type_args) == 1: |
| 41 | + item_type = var_type_args[0] |
| 42 | + setattr(obj, var_name, [item_type.from_dict(x) for x in data[var_name]]) |
| 43 | + else: |
| 44 | + setattr(obj, var_name, data.get(var_name, None)) |
| 45 | + return obj |
| 46 | + |
| 47 | + |
| 48 | + |
| 49 | +def fill_dict_from_object_using_type_hints(cls, obj): |
| 50 | + raise NotImplementedError |
| 51 | + |
| 52 | + |
| 53 | + |
| 54 | + |
| 55 | + |
| 56 | +######################################################################################## |
| 57 | +# Scenario 1: |
| 58 | +# Model a single file with FileDictModel, which is a dict at the top level. |
| 59 | +# Each key-value item is modeled by a DictItemModel. |
| 60 | + |
| 61 | + |
| 62 | + |
| 63 | +class FileDictModel(ABC, Generic[T]): |
| 64 | + """ |
| 65 | + A file base refers to a file that is stored in the database. |
| 66 | + At the top level the file must contain a dictionary with strings as keys. |
| 67 | + """ |
| 68 | + |
| 69 | + __file__ = None |
| 70 | + |
| 71 | + @classmethod |
| 72 | + def _get_item_model(cls): |
| 73 | + for base in cls.__orig_bases__: |
| 74 | + for type_args in get_args(base): |
| 75 | + if issubclass(type_args, FileDictItemModel): |
| 76 | + return type_args |
| 77 | + raise AttributeError( |
| 78 | + "FileDictModel must specify a FileDictItemModel " |
| 79 | + "(e.g. Users(FileDictModel[User]))" |
| 80 | + ) |
| 81 | + |
| 82 | + |
| 83 | + @classmethod |
| 84 | + def get_at_key(cls, key) -> T: |
| 85 | + """ |
| 86 | + Gets an item by key. |
| 87 | + The data is partially read from the __file__. |
| 88 | + """ |
| 89 | + data = io_safe.partial_read(cls.__file__, key) |
| 90 | + res: T = cls._get_item_model().from_key_value(key, data) |
| 91 | + return res |
| 92 | + |
| 93 | + @classmethod |
| 94 | + def session_at_key(cls, key): |
| 95 | + return cls._get_item_model().session(key) |
| 96 | + |
| 97 | + @classmethod |
| 98 | + def get_all(cls) -> dict[str, T]: |
| 99 | + data = io_safe.read(cls.__file__) |
| 100 | + return {k: cls._get_item_model().from_key_value(k, v) for k, v in data.items()} |
| 101 | + |
| 102 | + @classmethod |
| 103 | + def session(cls): |
| 104 | + """ |
| 105 | + Enter a session with the file as (session, data) where data is a dict of |
| 106 | + <key>: <ORM model of value> pairs. |
| 107 | + """ |
| 108 | + def make_session_obj_from_dict(data): |
| 109 | + sess_obj = {} |
| 110 | + for k, v in data.items(): |
| 111 | + sess_obj[k] = cls._get_item_model().from_key_value(k, v) |
| 112 | + return sess_obj |
| 113 | + return SessionFileFull(cls.__file__, make_session_obj_from_dict) |
| 114 | + |
| 115 | + |
| 116 | + @classmethod |
| 117 | + def get_where(cls, where: callable[str, T]) -> dict[str, T]: |
| 118 | + """ |
| 119 | + Return a dictionary of all the items for which the where function returns True. |
| 120 | + Where takes the key and the value's model object as arguments. |
| 121 | + """ |
| 122 | + return {k: v for k, v in cls.get_all().items() if where(k, v)} |
| 123 | + |
| 124 | + |
| 125 | + |
| 126 | + |
| 127 | +class FileDictItemModel(ABC): |
| 128 | + __key__: str |
| 129 | + |
| 130 | + @classmethod |
| 131 | + def from_key_value(cls: Type[T2], key, value) -> T2: |
| 132 | + obj = fill_object_from_dict_using_type_hints(cls(), cls, value) |
| 133 | + obj.__key__ = key |
| 134 | + return obj |
| 135 | + |
| 136 | + @classmethod |
| 137 | + def session(cls, key): |
| 138 | + def partial_func(x): |
| 139 | + return cls.from_key_value(key, x) |
| 140 | + return SessionFileKey(cls.__file__, key, partial_func) |
| 141 | + |
| 142 | + |
| 143 | + |
| 144 | +class DictModel(ABC): |
| 145 | + |
| 146 | + @classmethod |
| 147 | + def from_dict(cls, data) -> DictModel: |
| 148 | + obj = cls() |
| 149 | + return fill_object_from_dict_using_type_hints(obj, cls, data) |
| 150 | + |
| 151 | + def to_dict(self) -> dict: |
| 152 | + res = {} |
| 153 | + for var_name in get_type_hints(self).keys(): |
| 154 | + if (value := getattr(self, var_name)) is not None: |
| 155 | + res[var_name] = value |
| 156 | + return res |
| 157 | + |
| 158 | + |
| 159 | + |
| 160 | +######################################################################################## |
| 161 | +# Scenario 2: |
| 162 | +# Add in a later version of DDB |
| 163 | +# A folder containing multiple files, each containing json. |
| 164 | + |
| 165 | +# class FolderBase(ABC): |
| 166 | +# __folder__ = None |
| 167 | +# __file_model__: FileInFolderModel = None |
| 168 | + |
| 169 | + |
| 170 | +# class FileInFolderModel(ABC): |
| 171 | + |
| 172 | +# @classmethod |
| 173 | +# def get_by_name(cls, file_name: str) -> FileInFolderModel: |
| 174 | +# data = io_safe.read(f"{cls.__folder__}/{file_name}") |
| 175 | +# return cls(**data) |
0 commit comments