|
21 | 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
22 | 22 | SOFTWARE. |
23 | 23 | """ |
24 | | - |
25 | | -from typing import NamedTuple |
26 | | - |
27 | | -import sqlalchemy |
28 | | -from sqlalchemy.orm.relationships import RelationshipProperty |
29 | | - |
30 | | -from sqlalchemyseed import class_registry |
31 | | -from sqlalchemyseed import validator |
32 | | - |
33 | | - |
34 | | -class Entity(NamedTuple): |
35 | | - instance: object |
36 | | - attr_name: str |
37 | | - |
38 | | - @property |
39 | | - def cls_attribute(self): |
40 | | - return getattr(self.instance.__class__, self.attr_name) |
41 | | - |
42 | | - @property |
43 | | - def ins_attribute(self): |
44 | | - return getattr(self.instance, self.attr_name) |
45 | | - |
46 | | - @ins_attribute.setter |
47 | | - def ins_attribute(self, value): |
48 | | - setattr(self.instance, self.attr_name, value) |
49 | | - |
50 | | - |
51 | | -# def instantiate_class(class_, filtered_kwargs: dict, key: validator.Key, session: sqlalchemy.orm.Session = None): |
52 | | -# if key is validator.Key.data(): |
53 | | -# return class_(**filtered_kwargs) |
54 | | -# |
55 | | -# if key is validator.Key.filter() and session is not None: |
56 | | -# return session.query(class_).filter_by(**filtered_kwargs).one() |
57 | | - |
58 | | - |
59 | | -def filter_kwargs(kwargs: dict, class_, ref_prefix): |
60 | | - return { |
61 | | - k: v for k, v in kwargs.items() |
62 | | - if not str(k).startswith(ref_prefix) and not isinstance(getattr(class_, str(k)).property, RelationshipProperty) |
63 | | - } |
64 | | - |
65 | | - |
66 | | -def set_parent_attr_value(instance, parent: Entity): |
67 | | - if isinstance(parent.cls_attribute.property, RelationshipProperty): |
68 | | - if parent.cls_attribute.property.uselist is True: |
69 | | - parent.ins_attribute.append(instance) |
70 | | - else: |
71 | | - parent.ins_attribute = instance |
72 | | - |
73 | | - |
74 | | -def iter_ref_attr(attrs, ref_prefix): |
75 | | - for attr_name, value in attrs.items(): |
76 | | - if str(attr_name).startswith(ref_prefix): |
77 | | - # remove prefix of attr_name |
78 | | - yield str(attr_name)[len(ref_prefix):], value |
79 | | - |
80 | | - |
81 | | -class Seeder: |
82 | | - __model_key = validator.Key.model() |
83 | | - __data_key = validator.Key.data() |
84 | | - |
85 | | - def __init__(self, session: sqlalchemy.orm.Session = None, ref_prefix="!"): |
86 | | - self.session = session |
87 | | - self._class_registry = class_registry.ClassRegistry() |
88 | | - self._instances = [] |
89 | | - self.ref_prefix = ref_prefix |
90 | | - |
91 | | - @property |
92 | | - def instances(self): |
93 | | - return tuple(self._instances) |
94 | | - |
95 | | - # def get_model_class(self, entity, parent: Entity): |
96 | | - # model_label = self.__model_key.label |
97 | | - # if model_label in entity: |
98 | | - # class_path = entity[model_label] |
99 | | - # return self._class_registry.register_class(class_path) |
100 | | - # # parent is not None |
101 | | - # if isinstance(parent.attribute.property, RelationshipProperty): |
102 | | - # return parent.attribute.mapper.class_ |
103 | | - # else: # parent.attribute is instance of ColumnProperty |
104 | | - # table_name = parent.attribute.foreign_keys[0].table.name |
105 | | - # class_ = next( |
106 | | - # (mapper.class_ |
107 | | - # for mapper in parent.instance.__class__.registry.mappers |
108 | | - # if mapper.class_.__tablename__ == table_name), |
109 | | - # errors.ClassNotFoundError( |
110 | | - # "A class with table name '{}' is not found in the mappers".format(table_name)), |
111 | | - # ) |
112 | | - # return class_ |
113 | | - |
114 | | - def get_model_class(self, entity, parent: Entity): |
115 | | - if self.__model_key in entity: |
116 | | - return self._class_registry.register_class(entity[self.__model_key]) |
117 | | - # parent is not None |
118 | | - if isinstance(parent.cls_attribute.property, RelationshipProperty): |
119 | | - return parent.cls_attribute.mapper.class_ |
120 | | - |
121 | | - def seed(self, entities, add_to_session=True): |
122 | | - validator.SchemaValidator.validate( |
123 | | - entities, ref_prefix=self.ref_prefix) |
124 | | - |
125 | | - self._pre_seed(entities) |
126 | | - |
127 | | - if add_to_session: |
128 | | - self.session.add_all(self.instances) |
129 | | - |
130 | | - def _pre_seed(self, entity, parent: Entity = None): |
131 | | - if isinstance(entity, dict): |
132 | | - self._seed(entity, parent) |
133 | | - else: # is list |
134 | | - for item in entity: |
135 | | - self._pre_seed(item, parent) |
136 | | - |
137 | | - def _seed(self, entity, parent: Entity = None): |
138 | | - class_ = self.get_model_class(entity, parent) |
139 | | - # source_key: validator.Key = next( |
140 | | - # (sk for sk in self.__source_keys if sk.label in entity), None) |
141 | | - # source_data = entity[source_key.label] |
142 | | - |
143 | | - kwargs = entity[self.__data_key] |
144 | | - |
145 | | - # kwargs is list |
146 | | - if isinstance(kwargs, list): |
147 | | - for kwargs_ in kwargs: |
148 | | - instance = self._setup_instance(class_, kwargs_, parent) |
149 | | - self._seed_children(instance, kwargs_) |
150 | | - return |
151 | | - |
152 | | - # kwargs is dict |
153 | | - # instantiate object |
154 | | - instance = self._setup_instance(class_, kwargs, parent) |
155 | | - self._seed_children(instance, kwargs) |
156 | | - |
157 | | - def _seed_children(self, instance, kwargs): |
158 | | - for attr_name, value in iter_ref_attr(kwargs, self.ref_prefix): |
159 | | - self._pre_seed(entity=value, parent=Entity(instance, attr_name)) |
160 | | - |
161 | | - def _setup_instance(self, class_, kwargs: dict, parent: Entity): |
162 | | - instance = class_(**filter_kwargs(kwargs, class_, self.ref_prefix)) |
163 | | - if parent is not None: |
164 | | - set_parent_attr_value(instance, parent) |
165 | | - else: |
166 | | - self._instances.append(instance) |
167 | | - return instance |
168 | | - |
169 | | - # def instantiate_class(self, class_, kwargs: dict, key: validator.Key): |
170 | | - # filtered_kwargs = { |
171 | | - # k: v |
172 | | - # for k, v in kwargs.items() |
173 | | - # if not k.startswith("!") |
174 | | - # and not isinstance(getattr(class_, k), RelationshipProperty) |
175 | | - # } |
176 | | - # |
177 | | - # if key is validator.Key.data(): |
178 | | - # return class_(**filtered_kwargs) |
179 | | - # |
180 | | - # if key is validator.Key.filter() and self.session is not None: |
181 | | - # return self.session.query(class_).filter_by(**filtered_kwargs).one() |
182 | | - |
183 | | -# class HybridSeeder: |
184 | | -# __model_key = validator.Key.model() |
185 | | -# __source_keys = [validator.Key.data(), validator.Key.filter()] |
186 | | -# |
187 | | -# def __init__(self, session: sqlalchemy.orm.Session, ref_prefix): |
188 | | -# self.session = session |
189 | | -# self._class_registry = class_registry.ClassRegistry() |
190 | | -# self._instances = [] |
191 | | -# self.ref_prefix = ref_prefix |
192 | | -# |
193 | | -# @property |
194 | | -# def instances(self): |
195 | | -# return tuple(self._instances) |
196 | | -# |
197 | | -# def seed(self, entities): |
198 | | -# validator.SchemaValidator.validate( |
199 | | -# entities, ref_prefix=self.ref_prefix) |
200 | | -# |
201 | | -# self._pre_seed(entities) |
202 | | -# |
203 | | -# def _pre_seed(self, entity, parent=None): |
204 | | -# if isinstance(entity, dict): |
205 | | -# self._seed(entity, parent) |
206 | | -# else: # is list |
207 | | -# for item in entity: |
208 | | -# self._pre_seed(item, parent) |
209 | | -# |
210 | | -# def _seed(self, entity, parent): |
211 | | -# pass |
0 commit comments