root/django/branches/unicode/django/db/models/query.py

from django.conf import settings
from django.db import backend, connection, transaction
from django.db.models.fields import DateField, FieldDoesNotExist
from django.db.models import signals, loading
from django.dispatch import dispatcher
from django.utils.datastructures import SortedDict
from django.utils.encoding import smart_unicode
from django.contrib.contenttypes import generic
import datetime
import operator
import re

try:
    set
except NameError:
    from sets import Set as set   # Python 2.3 fallback

# The string constant used to separate query parts
LOOKUP_SEPARATOR = '__'

# The list of valid query types
QUERY_TERMS = (
    'exact', 'iexact', 'contains', 'icontains',
    'gt', 'gte', 'lt', 'lte', 'in',
    'startswith', 'istartswith', 'endswith', 'iendswith',
    'range', 'year', 'month', 'day', 'isnull', 'search',
    'regex', 'iregex',
)

# Size of each "chunk" for get_iterator calls.
# Larger values are slightly faster at the expense of more storage space.
GET_ITERATOR_CHUNK_SIZE = 100

class EmptyResultSet(Exception):
    pass

####################
# HELPER FUNCTIONS #
####################

# Django currently supports two forms of ordering.
# Form 1 (deprecated) example:
#     order_by=(('pub_date', 'DESC'), ('headline', 'ASC'), (None, 'RANDOM'))
# Form 2 (new-style) example:
#     order_by=('-pub_date', 'headline', '?')
# Form 1 is deprecated and will no longer be supported for Django's first
# official release. The following code converts from Form 1 to Form 2.

LEGACY_ORDERING_MAPPING = {'ASC': '_', 'DESC': '-_', 'RANDOM': '?'}

def handle_legacy_orderlist(order_list):
    if not order_list or isinstance(order_list[0], basestring):
        return order_list
    else:
        import warnings
        new_order_list = [LEGACY_ORDERING_MAPPING[j.upper()].replace('_', smart_unicode(i)) for i, j in order_list]
        warnings.warn("%r ordering syntax is deprecated. Use %r instead." % (order_list, new_order_list), DeprecationWarning)
        return new_order_list

def orderfield2column(f, opts):
    try:
        return opts.get_field(f, False).column
    except FieldDoesNotExist:
        return f

def orderlist2sql(order_list, opts, prefix=''):
    if prefix.endswith('.'):
        prefix = backend.quote_name(prefix[:-1]) + '.'
    output = []
    for f in handle_legacy_orderlist(order_list):
        if f.startswith('-'):
            output.append('%s%s DESC' % (prefix, backend.quote_name(orderfield2column(f[1:], opts))))
        elif f == '?':
            output.append(backend.get_random_function_sql())
        else:
            output.append('%s%s ASC' % (prefix, backend.quote_name(orderfield2column(f, opts))))
    return ', '.join(output)

def quote_only_if_word(word):
    if re.search('\W', word): # Don't quote if there are spaces or non-word chars.
        return word
    else:
        return backend.quote_name(word)

class _QuerySet(object):
    "Represents a lazy database lookup for a set of objects"
    def __init__(self, model=None):
        self.model = model
        self._filters = Q()
        self._order_by = None        # Ordering, e.g. ('date', '-name'). If None, use model's ordering.
        self._select_related = False # Whether to fill cache for related objects.
        self._max_related_depth = 0  # Maximum "depth" for select_related
        self._distinct = False       # Whether the query should use SELECT DISTINCT.
        self._select = {}            # Dictionary of attname -> SQL.
        self._where = []             # List of extra WHERE clauses to use.
        self._params = []            # List of params to use for extra WHERE clauses.
        self._tables = []            # List of extra tables to use.
        self._offset = None          # OFFSET clause.
        self._limit = None           # LIMIT clause.
        self._result_cache = None

    ########################
    # PYTHON MAGIC METHODS #
    ########################

    def __repr__(self):
        return repr(self._get_data())

    def __len__(self):
        return len(self._get_data())

    def __iter__(self):
        return iter(self._get_data())

    def __getitem__(self, k):
        "Retrieve an item or slice from the set of results."
        if not isinstance(k, (slice, int)):
            raise TypeError
        assert (not isinstance(k, slice) and (k >= 0)) \
            or (isinstance(k, slice) and (k.start is None or k.start >= 0) and (k.stop is None or k.stop >= 0)), \
            "Negative indexing is not supported."
        if self._result_cache is None:
            if isinstance(k, slice):
                # Offset:
                if self._offset is None:
                    offset = k.start
                elif k.start is None:
                    offset = self._offset
                else:
                    offset = self._offset + k.start
                # Now adjust offset to the bounds of any existing limit:
                if self._limit is not None and k.start is not None:
                    limit = self._limit - k.start
                else:
                    limit = self._limit

                # Limit:
                if k.stop is not None and k.start is not None:
                    if limit is None:
                        limit = k.stop - k.start
                    else:
                        limit = min((k.stop - k.start), limit)
                else:
                    if limit is None:
                        limit = k.stop
                    else:
                        if k.stop is not None:
                            limit = min(k.stop, limit)

                if k.step is None:
                    return self._clone(_offset=offset, _limit=limit)
                else:
                    return list(self._clone(_offset=offset, _limit=limit))[::k.step]
            else:
                try:
                    return list(self._clone(_offset=k, _limit=1))[0]
                except self.model.DoesNotExist, e:
                    raise IndexError, e.args
        else:
            return self._result_cache[k]

    def __and__(self, other):
        combined = self._combine(other)
        combined._filters = self._filters & other._filters
        return combined

    def __or__(self, other):
        combined = self._combine(other)
        combined._filters = self._filters | other._filters
        return combined

    ####################################
    # METHODS THAT DO DATABASE QUERIES #
    ####################################

    def iterator(self):
        "Performs the SELECT database lookup of this QuerySet."
        try:
            select, sql, params = self._get_sql_clause()
        except EmptyResultSet:
            raise StopIteration

        # self._select is a dictionary, and dictionaries' key order is
        # undefined, so we convert it to a list of tuples.
        extra_select = self._select.items()

        cursor = connection.cursor()
        cursor.execute("SELECT " + (self._distinct and "DISTINCT " or "") + ",".join(select) + sql, params)

        fill_cache = self._select_related
        fields = self.model._meta.fields
        index_end = len(fields)
        has_resolve_columns = hasattr(self, 'resolve_columns')
        while 1:
            rows = cursor.fetchmany(GET_ITERATOR_CHUNK_SIZE)
            if not rows:
                raise StopIteration
            for row in rows:
                if has_resolve_columns:
                    row = self.resolve_columns(row, fields)
                if fill_cache:
                    obj, index_end = get_cached_row(klass=self.model, row=row,
                                                    index_start=0, max_depth=self._max_related_depth)
                else:
                    obj = self.model(*row[:index_end])
                for i, k in enumerate(extra_select):
                    setattr(obj, k[0], row[index_end+i])
                yield obj

    def count(self):
        """
        Performs a SELECT COUNT() and returns the number of records as an
        integer.

        If the queryset is already cached (i.e. self._result_cache is set) this
        simply returns the length of the cached results set to avoid multiple
        SELECT COUNT(*) calls.
        """
        if self._result_cache is not None:
            return len(self._result_cache)

        counter = self._clone()
        counter._order_by = ()
        counter._select_related = False

        offset = counter._offset
        limit = counter._limit
        counter._offset = None
        counter._limit = None

        try:
            select, sql, params = counter._get_sql_clause()
        except EmptyResultSet:
            return 0

        cursor = connection.cursor()
        if self._distinct:
            id_col = "%s.%s" % (backend.quote_name(self.model._meta.db_table),
                    backend.quote_name(self.model._meta.pk.column))
            cursor.execute("SELECT COUNT(DISTINCT(%s))" % id_col + sql, params)
        else:
            cursor.execute("SELECT COUNT(*)" + sql, params)
        count = cursor.fetchone()[0]

        # Apply any offset and limit constraints manually, since using LIMIT or
        # OFFSET in SQL doesn't change the output of COUNT.
        if offset:
            count = max(0, count - offset)
        if limit:
            count = min(limit, count)

        return count

    def get(self, *args, **kwargs):
        "Performs the SELECT and returns a single object matching the given keyword arguments."
        clone = self.filter(*args, **kwargs)
        # clean up SQL by removing unneeded ORDER BY
        if not clone._order_by:
            clone._order_by = ()
        obj_list = list(clone)
        if len(obj_list) < 1:
            raise self.model.DoesNotExist, "%s matching query does not exist." % self.model._meta.object_name
        assert len(obj_list) == 1, "get() returned more than one %s -- it returned %s! Lookup parameters were %s" % (self.model._meta.object_name, len(obj_list), kwargs)
        return obj_list[0]

    def create(self, **kwargs):
        """
        Create a new object with the given kwargs, saving it to the database
        and returning the created object.
        """
        obj = self.model(**kwargs)
        obj.save()
        return obj

    def get_or_create(self, **kwargs):
        """
        Looks up an object with the given kwargs, creating one if necessary.
        Returns a tuple of (object, created), where created is a boolean
        specifying whether an object was created.
        """
        assert len(kwargs), 'get_or_create() must be passed at least one keyword argument'
        defaults = kwargs.pop('defaults', {})
        try:
            return self.get(**kwargs), False
        except self.model.DoesNotExist:
            params = dict([(k, v) for k, v in kwargs.items() if '__' not in k])
            params.update(defaults)
            obj = self.model(**params)
            obj.save()
            return obj, True

    def latest(self, field_name=None):
        """
        Returns the latest object, according to the model's 'get_latest_by'
        option or optional given field_name.
        """
        latest_by = field_name or self.model._meta.get_latest_by
        assert bool(latest_by), "latest() requires either a field_name parameter or 'get_latest_by' in the model"
        assert self._limit is None and self._offset is None, \
                "Cannot change a query once a slice has been taken."
        return self._clone(_limit=1, _order_by=('-'+latest_by,)).get()

    def in_bulk(self, id_list):
        """
        Returns a dictionary mapping each of the given IDs to the object with
        that ID.
        """
        assert self._limit is None and self._offset is None, \
                "Cannot use 'limit' or 'offset' with in_bulk"
        assert isinstance(id_list, (tuple,  list)), "in_bulk() must be provided with a list of IDs."
        id_list = list(id_list)
        if id_list == []:
            return {}
        qs = self._clone()
        qs._where.append("%s.%s IN (%s)" % (backend.quote_name(self.model._meta.db_table), backend.quote_name(self.model._meta.pk.column), ",".join(['%s'] * len(id_list))))
        qs._params.extend(id_list)
        return dict([(obj._get_pk_val(), obj) for obj in qs.iterator()])

    def delete(self):
        """
        Deletes the records in the current QuerySet.
        """
        assert self._limit is None and self._offset is None, \
            "Cannot use 'limit' or 'offset' with delete."

        del_query = self._clone()

        # disable non-supported fields
        del_query._select_related = False
        del_query._order_by = []

        # Delete objects in chunks to prevent an the list of
        # related objects from becoming too long
        more_objects = True
        while more_objects:
            # Collect all the objects to be deleted in this chunk, and all the objects
            # that are related to the objects that are to be deleted
            seen_objs = SortedDict()
            more_objects = False
            for object in del_query[0:GET_ITERATOR_CHUNK_SIZE]:
                more_objects = True
                object._collect_sub_objects(seen_objs)

            # If one or more objects were found, delete them.
            # Otherwise, stop looping.
            if more_objects:
                delete_objects(seen_objs)

        # Clear the result cache, in case this QuerySet gets reused.
        self._result_cache = None
    delete.alters_data = True

    ##################################################
    # PUBLIC METHODS THAT RETURN A QUERYSET SUBCLASS #
    ##################################################

    def values(self, *fields):
        return self._clone(klass=ValuesQuerySet, _fields=fields)

    def dates(self, field_name, kind, order='ASC'):
        """
        Returns a list of datetime objects representing all available dates
        for the given field_name, scoped to 'kind'.
        """
        assert kind in ("month", "year", "day"), "'kind' must be one of 'year', 'month' or 'day'."
        assert order in ('ASC', 'DESC'), "'order' must be either 'ASC' or 'DESC'."
        # Let the FieldDoesNotExist exception propagate.
        field = self.model._meta.get_field(field_name, many_to_many=False)
        assert isinstance(field, DateField), "%r isn't a DateField." % field_name
        return self._clone(klass=DateQuerySet, _field=field, _kind=kind, _order=order)

    ##################################################################
    # PUBLIC METHODS THAT ALTER ATTRIBUTES AND RETURN A NEW QUERYSET #
    ##################################################################

    def filter(self, *args, **kwargs):
        "Returns a new QuerySet instance with the args ANDed to the existing set."
        return self._filter_or_exclude(None, *args, **kwargs)

    def exclude(self, *args, **kwargs):
        "Returns a new QuerySet instance with NOT (args) ANDed to the existing set."
        return self._filter_or_exclude(QNot, *args, **kwargs)

    def _filter_or_exclude(self, mapper, *args, **kwargs):
        # mapper is a callable used to transform Q objects,
        # or None for identity transform
        if mapper is None:
            mapper = lambda x: x
        if len(args) > 0 or len(kwargs) > 0:
            assert self._limit is None and self._offset is None, \
                "Cannot filter a query once a slice has been taken."

        clone = self._clone()
        if len(kwargs) > 0:
            clone._filters = clone._filters & mapper(Q(**kwargs))
        if len(args) > 0:
            clone._filters = clone._filters & reduce(operator.and_, map(mapper, args))
        return clone

    def complex_filter(self, filter_obj):
        """Returns a new QuerySet instance with filter_obj added to the filters.
        filter_obj can be a Q object (has 'get_sql' method) or a dictionary of
        keyword lookup arguments."""
        # This exists to support framework features such as 'limit_choices_to',
        # and usually it will be more natural to use other methods.
        if hasattr(filter_obj, 'get_sql'):
            return self._filter_or_exclude(None, filter_obj)
        else:
            return self._filter_or_exclude(None, **filter_obj)

    def select_related(self, true_or_false=True, depth=0):
        "Returns a new QuerySet instance with '_select_related' modified."
        return self._clone(_select_related=true_or_false, _max_related_depth=depth)

    def order_by(self, *field_names):
        "Returns a new QuerySet instance with the ordering changed."
        assert self._limit is None and self._offset is None, \
                "Cannot reorder a query once a slice has been taken."
        return self._clone(_order_by=field_names)

    def distinct(self, true_or_false=True):
        "Returns a new QuerySet instance with '_distinct' modified."
        return self._clone(_distinct=true_or_false)

    def extra(self, select=None, where=None, params=None, tables=None):
        assert self._limit is None and self._offset is None, \
                "Cannot change a query once a slice has been taken"
        clone = self._clone()
        if select: clone._select.update(select)
        if where: clone._where.extend(where)
        if params: clone._params.extend(params)
        if tables: clone._tables.extend(tables)
        return clone

    ###################
    # PRIVATE METHODS #
    ###################

    def _clone(self, klass=None, **kwargs):
        if klass is None:
            klass = self.__class__
        c = klass()
        c.model = self.model
        c._filters = self._filters
        c._order_by = self._order_by
        c._select_related = self._select_related
        c._max_related_depth = self._max_related_depth
        c._distinct = self._distinct
        c._select = self._select.copy()
        c._where = self._where[:]
        c._params = self._params[:]
        c._tables = self._tables[:]
        c._offset = self._offset
        c._limit = self._limit
        c.__dict__.update(kwargs)
        return c

    def _combine(self, other):
        assert self._limit is None and self._offset is None \
            and other._limit is None and other._offset is None, \
            "Cannot combine queries once a slice has been taken."
        assert self._distinct == other._distinct, \
            "Cannot combine a unique query with a non-unique query"
        #  use 'other's order by
        #  (so that A.filter(args1) & A.filter(args2) does the same as
        #   A.filter(args1).filter(args2)
        combined = other._clone()
        if self._select: combined._select.update(self._select)
        if self._where: combined._where.extend(self._where)
        if self._params: combined._params.extend(self._params)
        if self._tables: combined._tables.extend(self._tables)
        # If 'self' is ordered and 'other' isn't, propagate 'self's ordering
        if (self._order_by is not None and len(self._order_by) > 0) and \
           (combined._order_by is None or len(combined._order_by) == 0):
            combined._order_by = self._order_by
        return combined

    def _get_data(self):
        if self._result_cache is None:
            self._result_cache = list(self.iterator())
        return self._result_cache

    def _get_sql_clause(self):
        opts = self.model._meta

        # Construct the fundamental parts of the query: SELECT X FROM Y WHERE Z.
        select = ["%s.%s" % (backend.quote_name(opts.db_table), backend.quote_name(f.column)) for f in opts