Index: django/db/models/sql/query.py =================================================================== --- django/db/models/sql/query.py (revision 7436) +++ django/db/models/sql/query.py (working copy) @@ -56,6 +56,7 @@ self.start_meta = None # SQL-related attributes + self.aggregates = [] self.select = [] self.tables = [] # Aliases in the order they are created. self.where = where() @@ -141,6 +142,7 @@ obj.standard_ordering = self.standard_ordering obj.start_meta = self.start_meta obj.select = self.select[:] + obj.aggregates = self.aggregates[:] obj.tables = self.tables[:] obj.where = deepcopy(self.where) obj.where_class = self.where_class @@ -174,6 +176,30 @@ row = self.resolve_columns(row, fields) yield row + def get_aggregation(self): + for field in self.select: + self.group_by.append(field) + self.select.extend(self.aggregates) + self.aggregates = [] + #print self.as_sql() + #print 'after', self.select + + get_name = lambda x : isinstance(x, tuple) and x[1] or x.aliased_name + + print 'final query', self.as_sql() + + if self.group_by: + data = self.execute_sql(MULTI) + result = [] + for rs in data.next(): + result.append(dict(zip([get_name(i) for i in self.select], rs))) + else: + data = self.execute_sql(SINGLE) + result = dict(zip([get_name(i) for i in self.select], data)) + + self.select = [] + return result + def get_count(self): """ Performs a COUNT() query using the current filter constraints. @@ -811,6 +837,81 @@ self.fill_related_selections(f.rel.to._meta, alias, cur_depth + 1, used, next, restricted) + def annotate(self, aggregate_expr, aliased_name, model): + field_list = aggregate_expr.split(LOOKUP_SEP) + opts = model._meta + + aggregate_func = field_list.pop() + + if len(field_list) > 1: + field, target, opts, join_list, last = self.setup_joins( + field_list, opts, self.get_initial_alias(), False) + final = len(join_list) + penultimate = last.pop() + if penultimate == final: + penultimate = last.pop() + if len(join_list) > 1: + extra = join_list[penultimate:] + final = penultimate + col = self.alias_map[extra[0]][LHS_JOIN_COL] + else: + col = target.column + + field_name = field_list.pop() + alias = join_list[-1] + alias = extra[final] + else: + field_name = field_list[0] + alias = opts.db_table + + + def add_aggregate(self, aggregate_expr, aliased_name, model): + """ + Adds a single aggregate expression to the Query + """ + + field_list = aggregate_expr.split(LOOKUP_SEP) + opts = model._meta + + aggregate_func = field_list.pop() + + if len(field_list) > 1: + field, target, opts, join_list, last = self.setup_joins( + field_list, opts, self.get_initial_alias(), False) + final = len(join_list) + penultimate = last.pop() + if penultimate == final: + penultimate = last.pop() + if len(join_list) > 1: + extra = join_list[penultimate:] + final = penultimate + col = self.alias_map[extra[0]][LHS_JOIN_COL] + else: + col = target.column + + field_name = field_list.pop() + alias = join_list[-1] + alias = extra[final] + else: + field_name = field_list[0] + alias = opts.db_table + + class AggregateNode: + def __init__(self, field_name, aggregate_func, aliased_name, alias): + self.field_name = field_name + self.aggregate_func = aggregate_func + self.aliased_name = aliased_name + self.alias = alias + + def as_sql(self, quote_func=None): + if not quote_func: + quote_func = lambda x: x + return '%s(%s.%s)' % (self.aggregate_func.upper(), + quote_func(self.alias), + quote_func(self.field_name)) + + self.aggregates.append(AggregateNode(field_name, aggregate_func, aliased_name, alias)) + def add_filter(self, filter_expr, connector=AND, negate=False, trim=False, single_filter=False): """ @@ -829,6 +930,10 @@ if not parts: raise FieldError("Cannot parse keyword query %r" % arg) + # if arg in (x.aliased_name for x in self.aggregates): + # self.having.append(arg) + # return + # Work out the lookup type and remove it from 'parts', if necessary. if len(parts) == 1 or parts[-1] not in self.query_terms: lookup_type = 'exact' Index: django/db/models/query.py =================================================================== --- django/db/models/query.py (revision 7436) +++ django/db/models/query.py (working copy) @@ -158,6 +158,34 @@ setattr(obj, k, row[i]) yield obj + def aggregate(self, *args, **kwargs): + """ + Returns the aggregation over the current model as a + dictionary. + + When applied to a ValuesQuerySet the results are GROUP BY-ed + by the fields specified in the values queryset. + + The kwargs are parsed as expression='alias'. + + If args is present the expression is passed as a kwarg with + itself as an alias. + """ + #Bug (or is it?): when doing both an aggregation on a related + #field and one on a 'local' field the local one goes + #wrong. something similar to: SELECT SUM(a.f1) FROM a INNER JOIN b; + #the value gets aggregated more than one time. + + if args: + newargs = {} + for arg in args: + newargs[arg] = arg + kwargs.update(newargs) + + for (aggregate_expr, alias) in kwargs.items(): + self.query.add_aggregate(aggregate_expr, alias, self.model) + return self.query.get_aggregation() + def count(self): """ Performs a SELECT COUNT() and returns the number of records as an @@ -326,6 +354,32 @@ """ return self._clone(klass=EmptyQuerySet) + def annotate(self, *args, **kwargs): + # Fix: Values is not working propperly + # Suffers from the same bug as aggrgate + # To-Do: HAVING + + if args: + newargs = {} + for arg in args: + newargs[arg] = arg + kwargs.update(newargs) + + opts = self.model._meta + fields = [] + + if isinstance(self, ValuesQuerySet): + obj = self._clone() + else: + fields.extend([f.name for f in opts.fields]) + obj = self._clone(klass=ValuesQuerySet, setup=True, _fields=fields) + + + for (aggregate_expr, alias) in kwargs.items(): + obj.query.add_aggregate(aggregate_expr, alias, self.model) + + return obj + ################################################################## # PUBLIC METHODS THAT ALTER ATTRIBUTES AND RETURN A NEW QUERYSET # ################################################################## @@ -335,7 +389,7 @@ Returns a new QuerySet that is a copy of the current one. This allows a QuerySet to proxy for a model manager in some cases. """ - return self._clone() + return self._clone() def filter(self, *args, **kwargs): """ @@ -488,14 +542,21 @@ # names of the model fields to select. def __iter__(self): + if self.query.aggregates: + return self.aggregate_iterator() return self.iterator() + def aggregate_iterator(self): + #Not lazy.. review + for i in self.query.get_aggregation(): + yield i + def iterator(self): self.query.trim_extra_select(self.extra_names) names = self.query.extra_select.keys() + self.field_names for row in self.query.results_iter(): yield dict(zip(names, row)) - + def _setup_query(self): """ Constructs the field_names list that the values query will be