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

1from functools import wraps 

2from uuid import UUID 

3 

4from pydantic import BaseModel, ConfigDict, Field 

5from sqlalchemy import select 

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

7from sqlalchemy.orm.mapper import Mapper 

8from sqlalchemy.orm.relationships import RelationshipProperty 

9from sqlalchemy.sql.base import ColumnCollection 

10 

11from .._model_base import SqlAlchemyBase 

12from .helpers import safe_call 

13 

14 

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

16 return {"id"} 

17 

18 

19class AutoInitConfig(BaseModel): 

20 """ 

21 Config class for `auto_init` decorator. 

22 """ 

23 

24 get_attr: str | None = None 

25 exclude: set = Field(default_factory=_default_exclusion) 

26 # auto_create: bool = False 

27 

28 

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

30 """ 

31 Returns the config for the given class. 

32 """ 

33 cfg = AutoInitConfig() 

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

35 # Get the config for the class 

36 try: 

37 class_config: ConfigDict = relation_cls.model_config 

38 except AttributeError: 

39 return cfg 

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

41 for attr in class_config: 

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

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

44 

45 return cfg 

46 

47 

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

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) 

59 

60 try: 

61 get_attr = cfg.get_attr 

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

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

64 except Exception: 

65 get_attr = "id" 

66 return get_attr 

67 

68 

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

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( 

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

78): 

79 elems_to_create: list[dict] = [] 

80 updated_elems: list[dict] = [] 

81 

82 cfg = _get_config(relation_cls) 

83 

84 for elem in all_elements: 

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

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

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

88 

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

90 elems_to_create.append(elem) 

91 continue 

92 

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

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

95 if key not in cfg.exclude: 

96 setattr(existing_elem, key, value) 

97 

98 updated_elems.append(existing_elem) 

99 

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

101 return new_elems + updated_elems 

102 

103 

104def auto_init(): # sourcery no-metrics 

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): 

113 @wraps(init) 

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

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__ 

124 config = _get_config(cls) 

125 exclude = config.exclude 

126 

127 alchemy_mapper: Mapper = self.__mapper__ 

128 model_columns: ColumnCollection = alchemy_mapper.columns 

129 relationships = alchemy_mapper.relationships 

130 

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

132 

133 if session is None: 

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

135 

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

137 if key in exclude: 

138 continue 

139 

140 if not hasattr(cls, key): 

141 continue 

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

143 

144 if key in model_columns: 

145 setattr(self, key, val) 

146 continue 

147 

148 if key in relationships: 

149 prop: RelationshipProperty = relationships[key] 

150 

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

152 relation_dir = prop.direction 

153 

154 # Identifies the parent class of the related object. 

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

156 

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

158 use_list: bool = prop.uselist 

159 

160 get_attr = get_lookup_attr(relation_cls) 

161 

162 if relation_dir == ONETOMANY and use_list: 

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

164 setattr(self, key, instances) 

165 

166 elif relation_dir == ONETOMANY: 

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

168 setattr(self, key, instance) 

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: 186 ↛ 136line 186 didn't jump to line 136 because the condition on line 186 was always true

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

188 setattr(self, key, instances) 

189 

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

191 

192 return wrapper 

193 

194 return decorator