Ticket #6811: functions.py

File functions.py, 6.8 KB (added by Justin, 17 years ago)

Expression/Function classes

Line 
1"""
2Classses for representing query functions and expressions. Not useful outside
3the SQL domain.
4
5Expressions do not need to inherit from E, but must implement the same methods
6"""
7
8class E(object):
9 """
10 Base class for expressions and functions.
11 Not abstract, can be used to give an alias to expressions resulting from
12 using Python operators, e.g. E(Sum('field1') + Sum('field2'), alias='a_name')
13 """
14 def __init__(self, expr, alias=None):
15 self.expr = expr
16 self._alias = alias
17
18 def __add__(self, rhs):
19 return Add(self, rhs)
20 def __sub__(self, rhs):
21 return Add(self, rhs)
22 def __mul__(self, rhs):
23 return Mul(self, rhs)
24 def __div__(self, rhs):
25 return Div(self, rhs)
26
27 def relabel_aliases(self, table_map, column_map=None):
28 """
29 Relabel the column alias, if necessary.
30 """
31 self.expr.relabel_aliases(table_map, column_map)
32
33 def as_sql(self, quote_func=None):
34 """
35 Returns the SQL string fragment for this object.
36
37 The quote_func function is used to quote the column components. If
38 None, it defaults to doing nothing.
39 """
40 return self.expr.as_sql(quote_func)
41
42 def get_cols(self):
43 """
44 Returns a list of columns used in the expression.
45 """
46 return self.expr.get_cols()
47
48 def output_alias(self):
49 """
50 Returns string to use as key in values dictionary
51 """
52 return self._alias
53
54 def do_group(self):
55 """
56 Returns whether this expression should trigger auto-grouping.
57 Only expressions containing Aggregates trigger auto-grouping.
58 """
59 return self.expr.do_group()
60
61
62class F(E):
63 """
64 Class to use a column in expressions, e.g:
65 F('price') < Avg('price')
66 F objects cannot contain arbitrary expressions, but only a column name.
67 """
68 def relabel_aliases(self, table_map, column_map=None):
69 c = self.expr
70 if isinstance(c, (list, tuple)):
71 self.expr = (table_map.get(c[0], c[0]), c[1])
72 elif isinstance(c, str):
73 if column_map:
74 col = column_map.get(c, c)
75 self.expr = (table_map.get('', ''), col)
76
77 def get_cols(self):
78 c = self.expr
79 if isinstance(c, (list, tuple)):
80 return [c[1]]
81 elif isinstance(c, str):
82 return [c]
83
84 def output_alias(self):
85 return self._alias or self.expr
86
87 def as_sql(self, quote_func=None):
88 if not quote_func:
89 quote_func = lambda x: x
90 expr = self.expr
91 if isinstance(expr, (list, tuple)):
92 expr = '%s.%s' % tuple([quote_func(c) for c in expr])
93 return '%s' % expr
94
95 def do_group(self):
96 return False
97
98
99class UnaryOp(E):
100 """
101 Base class for representing a unary operator such NOT
102 """
103 def as_sql(self, quote_func=None):
104 return '%s %s' % (self._op, self.expr.as_sql(quote_func))
105
106
107class BinaryOp(E):
108 """
109 Base class for representing a binary operator such as +, -, * or /
110 """
111 _op = None
112 def __init__(self, lhs, rhs, alias=None):
113 self.lhs = lhs
114 self.rhs = rhs
115 self._alias = alias
116
117 def relabel_aliases(self, table_map, column_map=None):
118 self.lhs.relabel_aliases(table_map, column_map)
119 self.rhs.relabel_aliases(table_map, column_map)
120
121 def get_cols(self):
122 return self.lhs.get_cols() + self.rhs.get_cols()
123
124 def as_sql(self, quote_func=None):
125 return '(%s %s %s)' % (self.lhs.as_sql(quote_func), self._op,
126 self.rhs.as_sql(quote_func))
127
128 def do_group(self):
129 return self.rhs.do_group() or self.lhs.do_group()
130
131
132class Add(BinaryOp):
133 _op = '+'
134
135class Sub(BinaryOp):
136 _op = '-'
137
138class Mul(BinaryOp):
139 _op = '*'
140
141class Div(BinaryOp):
142 _op = '/'
143
144class Not(UnaryOp):
145 _op = 'NOT'
146
147
148class Function(E):
149 """
150 Base class for query functions.
151 """
152 _func = None
153
154 def relabel_aliases(self, table_map, column_map=None):
155 c = self.expr
156 if isinstance(c, (list, tuple)):
157 self.expr = (table_map.get(c[0], c[0]), c[1])
158 elif isinstance(c, str):
159 if column_map:
160 col = column_map.get(c, c)
161 self.expr = (table_map.get('', ''), col)
162 else:
163 self.expr.relabel_aliases(table_map)
164
165 def get_cols(self):
166 c = self.expr
167 if isinstance(c, (list, tuple)):
168 return [c[1]]
169 elif isinstance(c, str):
170 return [c]
171 else:
172 return c.get_cols()
173
174 def output_alias(self):
175 return self._alias or (self.expr + '_' + self._func.lower())
176
177 def as_sql(self, quote_func=None):
178 if not quote_func:
179 quote_func = lambda x: x
180 expr = self.expr
181 if hasattr(expr, 'as_sql'):
182 expr = expr.as_sql(quote_func)
183 elif isinstance(expr, (list, tuple)):
184 expr = '%s.%s' % tuple([quote_func(c) for c in expr])
185 return '%s(%s)' % (self._func, expr)
186
187
188class Aggregate(Function):
189 """
190 Base class for query Aggregates.
191 An Aggregate in an expression will trigger auto-grouping.
192 Aggregates cannot contain arbitrary expressions, but only a column name.
193 """
194 #def __init__(self, col, alias=None):
195 # if not isinstance(c, (str, list, tuple)):
196 # raise FieldError("Invalid field name in aggregate function: '%s'" % col)
197
198 def do_group(self):
199 return False
200
201
202class Sum(Aggregate):
203 """
204 Perform a sum on the given column.
205 """
206 _func = 'SUM'
207
208class Avg(Aggregate):
209 """
210 Perform an average on the given column.
211 """
212 _func = 'AVG'
213
214class Min(Aggregate):
215 """
216 Select the minimum of the given column.
217 """
218 _func = 'MIN'
219
220class Max(Aggregate):
221 """
222 Select the maximum the given column.
223 """
224 _func = 'MAX'
225
226class Count(Aggregate):
227 _func = 'COUNT'
228 """
229 Perform a count on the given column.
230 """
231 def __init__(self, expr='*', alias=None, distinct=False):
232 self.distinct = distinct
233 super(Count, self).__init__(expr, alias=alias)
234
235 def get_cols(self):
236 if self.expr == '*':
237 return []
238 return super(Count, self).get_cols()
239
240 def relabel_aliases(self, table_map, column_map=None):
241 if self.expr == '*':
242 return
243 return super(Count, self).relabel_aliases(table_map, column_map)
244
245 def as_sql(self, quote_func=None):
246 if not quote_func:
247 quote_func = lambda x: x
248 expr = self.expr
249 if hasattr(expr, 'as_sql'):
250 expr = expr.as_sql(quote_func)
251 elif isinstance(expr, (list, tuple)):
252 expr = '%s.%s' % tuple([quote_func(c) for c in expr])
253 if self.distinct:
254 return 'COUNT(DISTINCT %s)' % expr
255 else:
256 return 'COUNT(%s)' % expr
257
258
Back to Top