Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy_utils/types/choice.py: 53%

1from enum import Enum 

2 

3from sqlalchemy import types 

4 

5from ..exceptions import ImproperlyConfigured 

6from .scalar_coercible import ScalarCoercible 

7 

8 

9class Choice: 

10 def __init__(self, code, value): 

11 self.code = code 

12 self.value = value 

13 

14 def __eq__(self, other): 

15 if isinstance(other, Choice): 

16 return self.code == other.code 

17 return other == self.code 

18 

19 def __hash__(self): 

20 return hash(self.code) 

21 

22 def __ne__(self, other): 

23 return not (self == other) 

24 

25 def __str__(self): 

26 return str(self.value) 

27 

28 def __repr__(self): 

29 return 'Choice(code={code}, value={value})'.format( 

30 code=self.code, value=self.value 

31 ) 

32 

33 

34class ChoiceType(ScalarCoercible, types.TypeDecorator): 

35 """ 

36 ChoiceType offers way of having fixed set of choices for given column. It 

37 could work with a list of tuple (a collection of key-value pairs), or 

38 integrate with :mod:`enum` in the standard library of Python 3. 

39 

40 Columns with ChoiceTypes are automatically coerced to Choice objects while 

41 a list of tuple been passed to the constructor. If a subclass of 

42 :class:`enum.Enum` is passed, columns will be coerced to :class:`enum.Enum` 

43 objects instead. 

44 

45 :: 

46 

47 class User(Base): 

48 TYPES = [ 

49 ('admin', 'Admin'), 

50 ('regular-user', 'Regular user') 

51 ] 

52 

53 __tablename__ = 'user' 

54 id = sa.Column(sa.Integer, primary_key=True) 

55 name = sa.Column(sa.Unicode(255)) 

56 type = sa.Column(ChoiceType(TYPES)) 

57 

58 

59 user = User(type='admin') 

60 user.type # Choice(code='admin', value='Admin') 

61 

62 Or:: 

63 

64 import enum 

65 

66 

67 class UserType(enum.Enum): 

68 admin = 1 

69 regular = 2 

70 

71 

72 class User(Base): 

73 __tablename__ = 'user' 

74 id = sa.Column(sa.Integer, primary_key=True) 

75 name = sa.Column(sa.Unicode(255)) 

76 type = sa.Column(ChoiceType(UserType, impl=sa.Integer())) 

77 

78 

79 user = User(type=1) 

80 user.type # <UserType.admin: 1> 

81 

82 

83 ChoiceType is very useful when the rendered values change based on user's 

84 locale: 

85 

86 :: 

87 

88 from babel import lazy_gettext as _ 

89 

90 

91 class User(Base): 

92 TYPES = [ 

93 ('admin', _('Admin')), 

94 ('regular-user', _('Regular user')) 

95 ] 

96 

97 __tablename__ = 'user' 

98 id = sa.Column(sa.Integer, primary_key=True) 

99 name = sa.Column(sa.Unicode(255)) 

100 type = sa.Column(ChoiceType(TYPES)) 

101 

102 

103 user = User(type='admin') 

104 user.type # Choice(code='admin', value='Admin') 

105 

106 print user.type # 'Admin' 

107 

108 Or:: 

109 

110 from enum import Enum 

111 from babel import lazy_gettext as _ 

112 

113 

114 class UserType(Enum): 

115 admin = 1 

116 regular = 2 

117 

118 

119 UserType.admin.label = _('Admin') 

120 UserType.regular.label = _('Regular user') 

121 

122 

123 class User(Base): 

124 __tablename__ = 'user' 

125 id = sa.Column(sa.Integer, primary_key=True) 

126 name = sa.Column(sa.Unicode(255)) 

127 type = sa.Column(ChoiceType(UserType, impl=sa.Integer())) 

128 

129 

130 user = User(type=UserType.admin) 

131 user.type # <UserType.admin: 1> 

132 

133 print user.type.label # 'Admin' 

134 """ 

135 

136 impl = types.Unicode(255) 

137 

138 cache_ok = True 

139 

140 def __init__(self, choices, impl=None): 

141 self.choices = tuple(choices) if isinstance(choices, list) else choices 

142 

143 if Enum is not None and isinstance(choices, type) and issubclass(choices, Enum): 

144 self.type_impl = EnumTypeImpl(enum_class=choices) 

145 else: 

146 self.type_impl = ChoiceTypeImpl(choices=choices) 

147 

148 if impl: 

149 self.impl = impl 

150 

151 @property 

152 def python_type(self): 

153 return self.impl.python_type 

154 

155 def _coerce(self, value): 

156 return self.type_impl._coerce(value) 

157 

158 def process_bind_param(self, value, dialect): 

159 return self.type_impl.process_bind_param(value, dialect) 

160 

161 def process_result_value(self, value, dialect): 

162 return self.type_impl.process_result_value(value, dialect) 

163 

164 

165class ChoiceTypeImpl: 

166 """The implementation for the ``Choice`` usage.""" 

167 

168 def __init__(self, choices): 

169 if not choices: 

170 raise ImproperlyConfigured('ChoiceType needs list of choices defined.') 

171 self.choices_dict = dict(choices) 

172 

173 def _coerce(self, value): 

174 if value is None: 

175 return value 

176 if isinstance(value, Choice): 

177 return value 

178 return Choice(value, self.choices_dict[value]) 

179 

180 def process_bind_param(self, value, dialect): 

181 if value and isinstance(value, Choice): 

182 return value.code 

183 return value 

184 

185 def process_result_value(self, value, dialect): 

186 if value: 

187 return Choice(value, self.choices_dict[value]) 

188 return value 

189 

190 

191class EnumTypeImpl: 

192 """The implementation for the ``Enum`` usage.""" 

193 

194 def __init__(self, enum_class): 

195 if not issubclass(enum_class, Enum): 

196 raise ImproperlyConfigured('EnumType needs a class of enum defined.') 

197 

198 self.enum_class = enum_class 

199 

200 def _coerce(self, value): 

201 if value is None: 

202 return None 

203 return self.enum_class(value) 

204 

205 def process_bind_param(self, value, dialect): 

206 if value is None: 

207 return None 

208 return self.enum_class(value).value 

209 

210 def process_result_value(self, value, dialect): 

211 return self._coerce(value)