Coverage for opt/mealie/lib/python3.12/site-packages/mealie/db/models/_model_utils/auto_init.py: 78%
104 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-11-25 15:32 +0000
1from functools import wraps 1j
2from uuid import UUID 1j
4from pydantic import BaseModel, ConfigDict, Field 1j
5from sqlalchemy import select 1j
6from sqlalchemy.orm import MANYTOMANY, MANYTOONE, ONETOMANY, Session 1j
7from sqlalchemy.orm.mapper import Mapper 1j
8from sqlalchemy.orm.relationships import RelationshipProperty 1j
9from sqlalchemy.sql.base import ColumnCollection 1j
11from .._model_base import SqlAlchemyBase 1j
12from .helpers import safe_call 1j
15def _default_exclusion() -> set[str]: 1j
16 return {"id"} 1jmlbcdefaghki
19class AutoInitConfig(BaseModel): 1j
20 """
21 Config class for `auto_init` decorator.
22 """
24 get_attr: str | None = None 1j
25 exclude: set = Field(default_factory=_default_exclusion) 1j
26 # auto_create: bool = False
29def _get_config(relation_cls: type[SqlAlchemyBase]) -> AutoInitConfig: 1j
30 """
31 Returns the config for the given class.
32 """
33 cfg = AutoInitConfig() 1jmlbcdefaghki
34 cfgKeys = cfg.model_dump().keys() 1jmlbcdefaghki
35 # Get the config for the class
36 try: 1jmlbcdefaghki
37 class_config: ConfigDict = relation_cls.model_config 1jmlbcdefaghki
38 except AttributeError: 1jmlbcdefaghki
39 return cfg 1jmlbcdefaghki
40 # Map all matching attributes in Config to all AutoInitConfig attributes
41 for attr in class_config: 1jmlbcdefaghki
42 if attr in cfgKeys: 42 ↛ 41line 42 didn't jump to line 41 because the condition on line 42 was always true1jmlbcdefaghki
43 setattr(cfg, attr, class_config[attr]) 1jmlbcdefaghki
45 return cfg 1jmlbcdefaghki
48def get_lookup_attr(relation_cls: type[SqlAlchemyBase]) -> str: 1j
49 """Returns the primary key attribute of the related class as a string.
51 Args:
52 relation_cls (DeclarativeMeta): The SQLAlchemy class to get the primary_key from
54 Returns:
55 Any: [description]
56 """
58 cfg = _get_config(relation_cls) 1lbcdefaghki
60 try: 1lbcdefaghki
61 get_attr = cfg.get_attr 1lbcdefaghki
62 if get_attr is None: 62 ↛ 66line 62 didn't jump to line 66 because the condition on line 62 was always true1lbcdefaghki
63 get_attr = relation_cls.__table__.primary_key.columns.keys()[0] 1lbcdefaghki
64 except Exception:
65 get_attr = "id"
66 return get_attr 1lbcdefaghki
69def handle_many_to_many(session, get_attr, relation_cls, all_elements: list[dict]): 1j
70 """
71 Proxy call to `handle_one_to_many_list` for many-to-many relationships. Because functionally, they do the same
72 """
73 return handle_one_to_many_list(session, get_attr, relation_cls, all_elements)
76def handle_one_to_many_list( 1j
77 session: Session, get_attr, relation_cls: type[SqlAlchemyBase], all_elements: list[dict] | list[str]
78):
79 elems_to_create: list[dict] = [] 1lbcdefaghki
80 updated_elems: list[dict] = [] 1lbcdefaghki
82 cfg = _get_config(relation_cls) 1lbcdefaghki
84 for elem in all_elements: 1lbcdefaghki
85 elem_id = elem.get(get_attr, None) if isinstance(elem, dict) else elem 1bcdefaghi
86 stmt = select(relation_cls).filter_by(**{get_attr: elem_id}) 1bcdefaghi
87 existing_elem = session.execute(stmt).scalars().one_or_none() 1bcdefaghi
89 if existing_elem is None and isinstance(elem, dict): 1bcdefaghi
90 elems_to_create.append(elem) 1bcdefaghi
91 continue 1bcdefaghi
93 elif isinstance(elem, dict): 93 ↛ 98line 93 didn't jump to line 98 because the condition on line 93 was always true1a
94 for key, value in elem.items(): 1a
95 if key not in cfg.exclude: 1a
96 setattr(existing_elem, key, value) 1a
98 updated_elems.append(existing_elem) 1a
100 new_elems = [safe_call(relation_cls, elem.copy(), session=session) for elem in elems_to_create] 1lbcdefaghki
101 return new_elems + updated_elems 1lbcdefaghki
104def auto_init(): # sourcery no-metrics 1j
105 """Wraps the `__init__` method of a class to automatically set the common
106 attributes.
108 Args:
109 exclude (Union[set, list], optional): [description]. Defaults to None.
110 """
112 def decorator(init): 1j
113 @wraps(init) 1j
114 def wrapper(self: SqlAlchemyBase, *args, **kwargs): # sourcery no-metrics 1j
115 """
116 Custom initializer that allows nested children initialization.
117 Only keys that are present as instance's class attributes are allowed.
118 These could be, for example, any mapped columns or relationships.
120 Code inspired from GitHub.
121 Ref: https://github.com/tiangolo/fastapi/issues/2194
122 """
123 cls = self.__class__ 1jmlbcdefaghki
124 config = _get_config(cls) 1jmlbcdefaghki
125 exclude = config.exclude 1jmlbcdefaghki
127 alchemy_mapper: Mapper = self.__mapper__ 1jmlbcdefaghki
128 model_columns: ColumnCollection = alchemy_mapper.columns 1jmlbcdefaghki
129 relationships = alchemy_mapper.relationships 1jmlbcdefaghki
131 session: Session = kwargs.get("session", None) 1jmlbcdefaghki
133 if session is None: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true1jmlbcdefaghki
134 raise ValueError("Session is required to initialize the model with `auto_init`")
136 for key, val in kwargs.items(): 1jmlbcdefaghki
137 if key in exclude: 1jmlbcdefaghki
138 continue 1jmlbcdefaghki
140 if not hasattr(cls, key): 1jmlbcdefaghki
141 continue 1jmlbcdefaghki
142 # raise TypeError(f"Invalid keyword argument: {key}")
144 if key in model_columns: 1jmlbcdefaghki
145 setattr(self, key, val) 1jmlbcdefaghki
146 continue 1jmlbcdefaghki
148 if key in relationships: 1lbcdefaghki
149 prop: RelationshipProperty = relationships[key] 1lbcdefaghki
151 # Identifies the type of relationship (ONETOMANY, MANYTOONE, many-to-one, many-to-many)
152 relation_dir = prop.direction 1lbcdefaghki
154 # Identifies the parent class of the related object.
155 relation_cls: type[SqlAlchemyBase] = prop.mapper.entity 1lbcdefaghki
157 # Identifies if the relationship was declared with use_list=True
158 use_list: bool = prop.uselist 1lbcdefaghki
160 get_attr = get_lookup_attr(relation_cls) 1lbcdefaghki
162 if relation_dir == ONETOMANY and use_list: 1lbcdefaghki
163 instances = handle_one_to_many_list(session, get_attr, relation_cls, val) 1lbcdefaghki
164 setattr(self, key, instances) 1lbcdefaghki
166 elif relation_dir == ONETOMANY: 166 ↛ 170line 166 didn't jump to line 170 because the condition on line 166 was always true1bcdefaghki
167 instance = safe_call(relation_cls, val.copy() if val else None, session=session) 1bcdefaghki
168 setattr(self, key, instance) 1bcdefaghki
170 elif relation_dir == MANYTOONE and not use_list:
171 if isinstance(val, dict):
172 val = val.get(get_attr)
174 if val is None:
175 raise ValueError(f"Expected 'id' to be provided for {key}")
177 if isinstance(val, str | int | UUID):
178 stmt = select(relation_cls).filter_by(**{get_attr: val})
179 instance = session.execute(stmt).scalars().one_or_none()
180 setattr(self, key, instance)
181 else:
182 # If the value is not of the type defined above we assume that it isn't a valid id
183 # and try a different approach.
184 pass
186 elif relation_dir == MANYTOMANY:
187 instances = handle_many_to_many(session, get_attr, relation_cls, val)
188 setattr(self, key, instances)
190 return init(self, *args, **kwargs) 1jmlbcdefaghki
192 return wrapper 1j
194 return decorator 1j