Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy_utils/generic.py: 28%

1import builtins 

2from collections.abc import Iterable 

3 

4import sqlalchemy as sa 

5from sqlalchemy.ext.hybrid import Comparator, hybrid_property 

6from sqlalchemy.orm import object_session 

7 

8from .exceptions import ImproperlyConfigured 

9from .functions import identity 

10 

11 

12class GenericComparator(Comparator): 

13 def __init__(self, relationship, cls): 

14 self.relationship = relationship 

15 self.cls = cls 

16 

17 def __clause_element__(self): 

18 raise NotImplementedError( 

19 'A generic relationship can only be compared against a model ' 

20 'instance or with is_type().' 

21 ) 

22 

23 def __eq__(self, other): 

24 discriminator, ids = self.relationship._attributes(self.cls) 

25 q = discriminator == type(other).__name__ 

26 for column, value in zip(ids, identity(other)): 

27 q &= column == value 

28 return q 

29 

30 def __ne__(self, other): 

31 return ~(self == other) 

32 

33 def is_type(self, other): 

34 discriminator, _ = self.relationship._attributes(self.cls) 

35 class_names = [ 

36 mapper.class_.__name__ for mapper in sa.inspect(other).self_and_descendants 

37 ] 

38 return discriminator.in_(class_names) 

39 

40 

41class GenericRelationship: 

42 """A generic form of the relationship property. 

43 

44 Creates a 1 to many relationship between the parent model 

45 and any other models using a discriminator (the class name). 

46 

47 Implemented as a hybrid property on top of the discriminator and id 

48 attributes, so it only relies on public SQLAlchemy APIs. 

49 

50 :param discriminator: 

51 Field to discriminate which model we are referring to. 

52 :param id: 

53 Field to point to the model we are referring to. 

54 """ 

55 

56 def __init__(self, discriminator, id): 

57 self.discriminator = discriminator 

58 if isinstance(id, Iterable) and not isinstance(id, str): 

59 self.ids = list(id) 

60 else: 

61 self.ids = [id] 

62 # Key for caching the assigned target object on the instance dict. 

63 self.cache_key = f'_generic_relationship_{builtins.id(self)}' 

64 

65 def _key(self, cls, column): 

66 if isinstance(column, str): 

67 return column 

68 if isinstance(column, hybrid_property): 

69 return column.__name__ 

70 try: 

71 return sa.inspect(cls).mapper.get_property_by_column(column).key 

72 except sa.orm.exc.UnmappedColumnError: 

73 raise ImproperlyConfigured( 

74 f'Could not find generic relationship attribute for {column!r}.' 

75 ) 

76 

77 def _keys(self, cls): 

78 return ( 

79 self._key(cls, self.discriminator), 

80 [self._key(cls, column) for column in self.ids], 

81 ) 

82 

83 def _attributes(self, cls): 

84 discriminator, ids = self._keys(cls) 

85 return getattr(cls, discriminator), [getattr(cls, key) for key in ids] 

86 

87 def _target_class(self, cls, discriminator): 

88 for mapper in sa.inspect(cls).registry.mappers: 

89 if mapper.class_.__name__ == discriminator: 

90 return mapper.class_ 

91 

92 def get(self, obj): 

93 discriminator_key, id_keys = self._keys(type(obj)) 

94 discriminator = getattr(obj, discriminator_key) 

95 id = tuple(getattr(obj, key) for key in id_keys) 

96 

97 # Return the object assigned to this relationship as long as the 

98 # discriminator and id still point to it. 

99 cached = obj.__dict__.get(self.cache_key) 

100 if ( 

101 cached is not None 

102 and type(cached).__name__ == discriminator 

103 and identity(cached) == id 

104 ): 

105 return cached 

106 

107 # Use the session the object is bound to in order to perform 

108 # a lazy query for the target. 

109 session = object_session(obj) 

110 if session is None: 

111 return None 

112 

113 target_class = self._target_class(type(obj), discriminator) 

114 if target_class is None: 

115 # Unknown discriminator; return nothing. 

116 return None 

117 

118 return session.get(target_class, id) 

119 

120 def set(self, obj, value): 

121 discriminator_key, id_keys = self._keys(type(obj)) 

122 obj.__dict__[self.cache_key] = value 

123 

124 if value is None: 

125 # Nullify relationship args 

126 for key in id_keys: 

127 setattr(obj, key, None) 

128 setattr(obj, discriminator_key, None) 

129 else: 

130 for key, pk in zip(id_keys, identity(value)): 

131 setattr(obj, key, pk) 

132 setattr(obj, discriminator_key, type(value).__name__) 

133 

134 def comparator(self, cls): 

135 return GenericComparator(self, cls) 

136 

137 

138def generic_relationship(discriminator, id, doc=None): 

139 relationship = GenericRelationship(discriminator, id) 

140 prop = hybrid_property( 

141 relationship.get, 

142 relationship.set, 

143 custom_comparator=relationship.comparator, 

144 ) 

145 prop.__doc__ = doc 

146 return prop