Coverage for opt/mealie/lib/python3.12/site-packages/mealie/db/models/_model_utils/auto_init.py: 78%

1from functools import wraps 1 ctx1l

2from uuid import UUID 1 ctx1l

3 

4from pydantic import BaseModel, ConfigDict, Field 1 ctx1l

5from sqlalchemy import select 1 ctx1l

6from sqlalchemy.orm import MANYTOMANY, MANYTOONE, ONETOMANY, Session 1 ctx1l

7from sqlalchemy.orm.mapper import Mapper 1 ctx1l

8from sqlalchemy.orm.relationships import RelationshipProperty 1 ctx1l

9from sqlalchemy.sql.base import ColumnCollection 1 ctx1l

10 

11from .._model_base import SqlAlchemyBase 1 ctx1l

12from .helpers import safe_call 1 ctx1l

13 

14 

15def _default_exclusion() -> set[str]: 1 ctx1l

16 return {"id"} 17 ctx1lpmqnbcdefghiajok

17 

18 

19class AutoInitConfig(BaseModel): 1 ctx1l

20 """ 

21 Config class for `auto_init` decorator. 

22 """ 

23 

24 get_attr: str | None = None 1 ctx1l

25 exclude: set = Field(default_factory=_default_exclusion) 1 ctx1l

26 # auto_create: bool = False 

27 

28 

29def _get_config(relation_cls: type[SqlAlchemyBase]) -> AutoInitConfig: 1 ctx1l

30 """ 

31 Returns the config for the given class. 

32 """ 

33 cfg = AutoInitConfig() 17 ctx1lpmqnbcdefghiajok

34 cfgKeys = cfg.model_dump().keys() 17 ctx1lpmqnbcdefghiajok

35 # Get the config for the class 

36 try: 17 ctx1lpmqnbcdefghiajok

37 class_config: ConfigDict = relation_cls.model_config 17 ctx1lpmqnbcdefghiajok

38 except AttributeError: 17 ctx1lpmqnbcdefghiajok

39 return cfg 17 ctx1lpmqnbcdefghiajok

40 # Map all matching attributes in Config to all AutoInitConfig attributes 

41 for attr in class_config: 17 ctx1lpmqnbcdefghiajok

42 if attr in cfgKeys: 42 ↛ 41line 42 didn't jump to line 41 because the condition on line 42 was always true17 ctx1lpmqnbcdefghiajok

43 setattr(cfg, attr, class_config[attr]) 17 ctx1lpmqnbcdefghiajok

44 

45 return cfg 17 ctx1lpmqnbcdefghiajok

46 

47 

48def get_lookup_attr(relation_cls: type[SqlAlchemyBase]) -> str: 1 ctx1l

49 """Returns the primary key attribute of the related class as a string. 

50 

51 Args: 

52 relation_cls (DeclarativeMeta): The SQLAlchemy class to get the primary_key from 

53 

54 Returns: 

55 Any: [description] 

56 """ 

57 

58 cfg = _get_config(relation_cls) 14 ctx1mnbcdefghiajok

59 

60 try: 14 ctx1mnbcdefghiajok

61 get_attr = cfg.get_attr 14 ctx1mnbcdefghiajok

62 if get_attr is None: 62 ↛ 66line 62 didn't jump to line 66 because the condition on line 62 was always true14 ctx1mnbcdefghiajok

63 get_attr = relation_cls.__table__.primary_key.columns.keys()[0] 14 ctx1mnbcdefghiajok

64 except Exception: 

65 get_attr = "id" 

66 return get_attr 14 ctx1mnbcdefghiajok

67 

68 

69def handle_many_to_many(session, get_attr, relation_cls, all_elements: list[dict]): 1 ctx1l

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) 

74 

75 

76def handle_one_to_many_list( 1 ctx1l

77 session: Session, get_attr, relation_cls: type[SqlAlchemyBase], all_elements: list[dict] | list[str] 

78): 

79 elems_to_create: list[dict] = [] 14 ctx1mnbcdefghiajok

80 updated_elems: list[dict] = [] 14 ctx1mnbcdefghiajok

81 

82 cfg = _get_config(relation_cls) 14 ctx1mnbcdefghiajok

83 

84 for elem in all_elements: 14 ctx1mnbcdefghiajok

85 elem_id = elem.get(get_attr, None) if isinstance(elem, dict) else elem 11 ctx1bcdefghiajk

86 stmt = select(relation_cls).filter_by(**{get_attr: elem_id}) 11 ctx1bcdefghiajk

87 existing_elem = session.execute(stmt).scalars().one_or_none() 11 ctx1bcdefghiajk

88 

89 if existing_elem is None and isinstance(elem, dict): 11 ctx1bcdefghiajk

90 elems_to_create.append(elem) 11 ctx1bcdefghiajk

91 continue 11 ctx1bcdefghiajk

92 

93 elif isinstance(elem, dict): 93 ↛ 98line 93 didn't jump to line 98 because the condition on line 93 was always true1 ctx1a

94 for key, value in elem.items(): 1 ctx1a

95 if key not in cfg.exclude: 1 ctx1a

96 setattr(existing_elem, key, value) 1 ctx1a

97 

98 updated_elems.append(existing_elem) 1 ctx1a

99 

100 new_elems = [safe_call(relation_cls, elem.copy(), session=session) for elem in elems_to_create] 14 ctx1mnbcdefghiajok

101 return new_elems + updated_elems 14 ctx1mnbcdefghiajok

102 

103 

104def auto_init(): # sourcery no-metrics 1 ctx1l

105 """Wraps the `__init__` method of a class to automatically set the common 

106 attributes. 

107 

108 Args: 

109 exclude (Union[set, list], optional): [description]. Defaults to None. 

110 """ 

111 

112 def decorator(init): 1 ctx1l

113 @wraps(init) 1 ctx1l

114 def wrapper(self: SqlAlchemyBase, *args, **kwargs): # sourcery no-metrics 1 ctx1l

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. 

119 

120 Code inspired from GitHub. 

121 Ref: https://github.com/tiangolo/fastapi/issues/2194 

122 """ 

123 cls = self.__class__ 17 ctx1lpmqnbcdefghiajok

124 config = _get_config(cls) 17 ctx1lpmqnbcdefghiajok

125 exclude = config.exclude 17 ctx1lpmqnbcdefghiajok

126 

127 alchemy_mapper: Mapper = self.__mapper__ 17 ctx1lpmqnbcdefghiajok

128 model_columns: ColumnCollection = alchemy_mapper.columns 17 ctx1lpmqnbcdefghiajok

129 relationships = alchemy_mapper.relationships 17 ctx1lpmqnbcdefghiajok

130 

131 session: Session = kwargs.get("session", None) 17 ctx1lpmqnbcdefghiajok

132 

133 if session is None: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true17 ctx1lpmqnbcdefghiajok

134 raise ValueError("Session is required to initialize the model with `auto_init`") 

135 

136 for key, val in kwargs.items(): 17 ctx1lpmqnbcdefghiajok

137 if key in exclude: 17 ctx1lpmqnbcdefghiajok

138 continue 17 ctx1lpmqnbcdefghiajok

139 

140 if not hasattr(cls, key): 17 ctx1lpmqnbcdefghiajok

141 continue 17 ctx1lpmqnbcdefghiajok

142 # raise TypeError(f"Invalid keyword argument: {key}") 

143 

144 if key in model_columns: 17 ctx1lpmqnbcdefghiajok

145 setattr(self, key, val) 17 ctx1lpmqnbcdefghiajok

146 continue 17 ctx1lpmqnbcdefghiajok

147 

148 if key in relationships: 14 ctx1mnbcdefghiajok

149 prop: RelationshipProperty = relationships[key] 14 ctx1mnbcdefghiajok

150 

151 # Identifies the type of relationship (ONETOMANY, MANYTOONE, many-to-one, many-to-many) 

152 relation_dir = prop.direction 14 ctx1mnbcdefghiajok

153 

154 # Identifies the parent class of the related object. 

155 relation_cls: type[SqlAlchemyBase] = prop.mapper.entity 14 ctx1mnbcdefghiajok

156 

157 # Identifies if the relationship was declared with use_list=True 

158 use_list: bool = prop.uselist 14 ctx1mnbcdefghiajok

159 

160 get_attr = get_lookup_attr(relation_cls) 14 ctx1mnbcdefghiajok

161 

162 if relation_dir == ONETOMANY and use_list: 14 ctx1mnbcdefghiajok

163 instances = handle_one_to_many_list(session, get_attr, relation_cls, val) 14 ctx1mnbcdefghiajok

164 setattr(self, key, instances) 14 ctx1mnbcdefghiajok

165 

166 elif relation_dir == ONETOMANY: 166 ↛ 170line 166 didn't jump to line 170 because the condition on line 166 was always true12 ctx1mbcdefghiajk

167 instance = safe_call(relation_cls, val.copy() if val else None, session=session) 12 ctx1mbcdefghiajk

168 setattr(self, key, instance) 12 ctx1mbcdefghiajk

169 

170 elif relation_dir == MANYTOONE and not use_list: 

171 if isinstance(val, dict): 

172 val = val.get(get_attr) 

173 

174 if val is None: 

175 raise ValueError(f"Expected 'id' to be provided for {key}") 

176 

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 

185 

186 elif relation_dir == MANYTOMANY: 

187 instances = handle_many_to_many(session, get_attr, relation_cls, val) 

188 setattr(self, key, instances) 

189 

190 return init(self, *args, **kwargs) 17 ctx1lpmqnbcdefghiajok

191 

192 return wrapper 1 ctx1l

193 

194 return decorator 1 ctx1l