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

1from functools import wraps 1 ctx1r

2from uuid import UUID 1 ctx1r

3 

4from pydantic import BaseModel, ConfigDict, Field 1 ctx1r

5from sqlalchemy import select 1 ctx1r

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

7from sqlalchemy.orm.mapper import Mapper 1 ctx1r

8from sqlalchemy.orm.relationships import RelationshipProperty 1 ctx1r

9from sqlalchemy.sql.base import ColumnCollection 1 ctx1r

10 

11from .._model_base import SqlAlchemyBase 1 ctx1r

12from .helpers import safe_call 1 ctx1r

13 

14 

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

16 return {"id"} 23 ctx1rustvdewfgbhijklmnocpqa

17 

18 

19class AutoInitConfig(BaseModel): 1 ctx1r

20 """ 

21 Config class for `auto_init` decorator. 

22 """ 

23 

24 get_attr: str | None = None 1 ctx1r

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

26 # auto_create: bool = False 

27 

28 

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

30 """ 

31 Returns the config for the given class. 

32 """ 

33 cfg = AutoInitConfig() 23 ctx1rustvdewfgbhijklmnocpqa

34 cfgKeys = cfg.model_dump().keys() 23 ctx1rustvdewfgbhijklmnocpqa

35 # Get the config for the class 

36 try: 23 ctx1rustvdewfgbhijklmnocpqa

37 class_config: ConfigDict = relation_cls.model_config 23 ctx1rustvdewfgbhijklmnocpqa

38 except AttributeError: 23 ctx1rustvdewfgbhijklmnocpqa

39 return cfg 23 ctx1rustvdewfgbhijklmnocpqa

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

41 for attr in class_config: 21 ctx1rustdefgbhijklmnocpqa

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

43 setattr(cfg, attr, class_config[attr]) 21 ctx1rustdefgbhijklmnocpqa

44 

45 return cfg 21 ctx1rustdefgbhijklmnocpqa

46 

47 

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

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) 19 ctx1stdefgbhijklmnocpqa

59 

60 try: 19 ctx1stdefgbhijklmnocpqa

61 get_attr = cfg.get_attr 19 ctx1stdefgbhijklmnocpqa

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

63 get_attr = relation_cls.__table__.primary_key.columns.keys()[0] 19 ctx1stdefgbhijklmnocpqa

64 except Exception: 

65 get_attr = "id" 

66 return get_attr 19 ctx1stdefgbhijklmnocpqa

67 

68 

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

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) 2 ctx1da

74 

75 

76def handle_one_to_many_list( 1 ctx1r

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

78): 

79 elems_to_create: list[dict] = [] 19 ctx1stdefgbhijklmnocpqa

80 updated_elems: list[dict] = [] 19 ctx1stdefgbhijklmnocpqa

81 

82 cfg = _get_config(relation_cls) 19 ctx1stdefgbhijklmnocpqa

83 

84 for elem in all_elements: 19 ctx1stdefgbhijklmnocpqa

85 elem_id = elem.get(get_attr, None) if isinstance(elem, dict) else elem 16 ctx1efgbhijklmnocpqa

86 stmt = select(relation_cls).filter_by(**{get_attr: elem_id}) 16 ctx1efgbhijklmnocpqa

87 existing_elem = session.execute(stmt).scalars().one_or_none() 16 ctx1efgbhijklmnocpqa

88 

89 if existing_elem is None and isinstance(elem, dict): 16 ctx1efgbhijklmnocpqa

90 elems_to_create.append(elem) 15 ctx1fgbhijklmnocpqa

91 continue 15 ctx1fgbhijklmnocpqa

92 

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

94 for key, value in elem.items(): 3 ctx1ebc

95 if key not in cfg.exclude: 3 ctx1ebc

96 setattr(existing_elem, key, value) 3 ctx1ebc

97 

98 updated_elems.append(existing_elem) 3 ctx1ebc

99 

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

101 return new_elems + updated_elems 19 ctx1stdefgbhijklmnocpqa

102 

103 

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

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 ctx1r

113 @wraps(init) 1 ctx1r

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

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__ 23 ctx1rustvdewfgbhijklmnocpqa

124 config = _get_config(cls) 23 ctx1rustvdewfgbhijklmnocpqa

125 exclude = config.exclude 23 ctx1rustvdewfgbhijklmnocpqa

126 

127 alchemy_mapper: Mapper = self.__mapper__ 23 ctx1rustvdewfgbhijklmnocpqa

128 model_columns: ColumnCollection = alchemy_mapper.columns 23 ctx1rustvdewfgbhijklmnocpqa

129 relationships = alchemy_mapper.relationships 23 ctx1rustvdewfgbhijklmnocpqa

130 

131 session: Session = kwargs.get("session", None) 23 ctx1rustvdewfgbhijklmnocpqa

132 

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

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

135 

136 for key, val in kwargs.items(): 23 ctx1rustvdewfgbhijklmnocpqa

137 if key in exclude: 23 ctx1rustvdewfgbhijklmnocpqa

138 continue 21 ctx1rustdefgbhijklmnocpqa

139 

140 if not hasattr(cls, key): 23 ctx1rustvdewfgbhijklmnocpqa

141 continue 23 ctx1rustvdewfgbhijklmnocpqa

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

143 

144 if key in model_columns: 23 ctx1rustvdewfgbhijklmnocpqa

145 setattr(self, key, val) 23 ctx1rustvdewfgbhijklmnocpqa

146 continue 23 ctx1rustvdewfgbhijklmnocpqa

147 

148 if key in relationships: 19 ctx1stdefgbhijklmnocpqa

149 prop: RelationshipProperty = relationships[key] 19 ctx1stdefgbhijklmnocpqa

150 

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

152 relation_dir = prop.direction 19 ctx1stdefgbhijklmnocpqa

153 

154 # Identifies the parent class of the related object. 

155 relation_cls: type[SqlAlchemyBase] = prop.mapper.entity 19 ctx1stdefgbhijklmnocpqa

156 

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

158 use_list: bool = prop.uselist 19 ctx1stdefgbhijklmnocpqa

159 

160 get_attr = get_lookup_attr(relation_cls) 19 ctx1stdefgbhijklmnocpqa

161 

162 if relation_dir == ONETOMANY and use_list: 19 ctx1stdefgbhijklmnocpqa

163 instances = handle_one_to_many_list(session, get_attr, relation_cls, val) 19 ctx1stdefgbhijklmnocpqa

164 setattr(self, key, instances) 19 ctx1stdefgbhijklmnocpqa

165 

166 elif relation_dir == ONETOMANY: 15 ctx1dfgbhijklmnocpa

167 instance = safe_call(relation_cls, val.copy() if val else None, session=session) 14 ctx1fgbhijklmnocpa

168 setattr(self, key, instance) 14 ctx1fgbhijklmnocpa

169 

170 elif relation_dir == MANYTOONE and not use_list: 2 ctx1da

171 if isinstance(val, dict): 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true2 ctx1da

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): 177 ↛ 178line 177 didn't jump to line 178 because the condition on line 177 was never true2 ctx1da

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 2 ctx1da

185 

186 elif relation_dir == MANYTOMANY: 186 ↛ 136line 186 didn't jump to line 136 because the condition on line 186 was always true2 ctx1da

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

188 setattr(self, key, instances) 2 ctx1da

189 

190 return init(self, *args, **kwargs) 23 ctx1rustvdewfgbhijklmnocpqa

191 

192 return wrapper 1 ctx1r

193 

194 return decorator 1 ctx1r