@@ -51,12 +51,26 @@ def __getattr__(self, item):
51
51
_none = object ()
52
52
53
53
54
+ def _get_column (model , column_or_name ) -> Column :
55
+ if isinstance (column_or_name , str ):
56
+ return getattr (model , column_or_name )
57
+
58
+ if isinstance (column_or_name , Column ):
59
+ if column_or_name in model :
60
+ return column_or_name
61
+ raise AttributeError ('Column {} does not belong to model {}' .format (
62
+ column_or_name , model ))
63
+
64
+ raise TypeError ('Unknown column {} with type {}' .
65
+ format (column_or_name , type (column_or_name )))
66
+
67
+
54
68
class ModelLoader (Loader ):
55
- def __init__ (self , model , * column_names , ** extras ):
69
+ def __init__ (self , model , * columns , ** extras ):
56
70
self .model = model
57
71
self ._distinct = None
58
- if column_names :
59
- self .columns = self . _column_loader (model , column_names )
72
+ if columns :
73
+ self .columns = [ _get_column (model , name ) for name in columns ]
60
74
else :
61
75
self .columns = model
62
76
self .extras = dict ((key , self .get (value ))
@@ -121,30 +135,14 @@ def get_from(self):
121
135
getattr (subloader , 'on_clause' , None ))
122
136
return rv
123
137
124
- def load (self , * column_names , ** extras ):
125
- if column_names :
126
- self .columns = self . _column_loader (self .model , column_names )
138
+ def load (self , * columns , ** extras ):
139
+ if columns :
140
+ self .columns = [ _get_column (self .model , name ) for name in columns ]
127
141
128
142
self .extras .update ((key , self .get (value ))
129
143
for key , value in extras .items ())
130
144
return self
131
145
132
- @classmethod
133
- def _column_loader (cls , model , column_names ):
134
- def column_formatter (column_name ):
135
- if isinstance (column_name , str ):
136
- return getattr (model , column_name )
137
- elif isinstance (column_name , Column ):
138
- if column_name not in model :
139
- raise AttributeError ('Column {} does not belong '
140
- 'to this model' .format (column_name ))
141
- return column_name
142
- else :
143
- raise TypeError ('Unknown column name {} type {}' .
144
- format (column_name , type (column_name )))
145
-
146
- return [column_formatter (column_name ) for column_name in column_names ]
147
-
148
146
def on (self , on_clause ):
149
147
self .on_clause = on_clause
150
148
return self
@@ -163,8 +161,8 @@ def none_as_none(self, enabled=True):
163
161
164
162
165
163
class AliasLoader (ModelLoader ):
166
- def __init__ (self , alias , * column_names , ** extras ):
167
- super ().__init__ (alias , * column_names , ** extras )
164
+ def __init__ (self , alias , * columns , ** extras ):
165
+ super ().__init__ (alias , * columns , ** extras )
168
166
169
167
170
168
class ColumnLoader (Loader ):
0 commit comments