from django.db import models
from django.db.models.query import QuerySet


class QuerySetExplainMixin:
  """Dump SQL EXPLAIN for a Queryset."""
  def explain(self, analyze=False, verbose=False):
    from django.db import connections
    cursor = connections[self.db].cursor()
    cmd = 'explain'
    cmd += ' analyze' if analyze else ''
    cmd += ' verbose' if verbose else ''
    cursor.execute('%s %s' % (cmd, str(self.query)))
    return cursor.fetchall()


class QuerySetFieldValuesMixin:
  """Returns generator of values for one field on a Queryset."""
  def field_values(self, field_name, distinct=False):
    values = self.values(field_name)
    if distinct:
      values = values.distinct()
    return (row[field_name] for row in values)


class QuerySetCountGroupbyMixin:
  """
  Useful for debugging in shell.

  Return a COUNT(*) GROUP BY breakdown of queryset by a specified attribute,
  i.e. a sorted (descending) list of dictionaries {attr value: count} of each
  represented attribute.

  By default, print human readable breakdown.
  """
  def count_groupby(self, attr, silent=False):
    counts = (self.values(attr).annotate(count=models.Count(attr))
              .order_by('count').reverse())

    choices = dict(getattr(self.model._meta.get_field(attr), 'choices', []))
    if not silent and choices.keys():
      for group in counts:
        print group['count'], '\t', choices.get(group[attr])

    return counts


QuerySet.__bases__ += (QuerySetExplainMixin, QuerySetFieldValuesMixin,
                       QuerySetCountGroupbyMixin)