#!/usr/bin/python
# -*- encoding: utf-8 -*-

##
## ambidjangolib/db.py
## Created on Wed Apr  5 23:26:18 2006
## by Antti Kaihola <akaihola@ambitone.com>
## 
## Copyright (C) 2006 Antti Kaihola
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 2 of the License, or
## (at your option) any later version.
## 
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
## 
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##

def preload_related(objects, model, *related):
    """
    Given a queryset of objects, retrieves related objects and stores
    them as lists into attributes of the original objects.

    See http://code.djangoproject.com/wiki/CookBookPreloadRelated for
    more information.

    Example:
    >>> users = User.objects.filter(username__startswith='a')
    >>> preload_related(users,
    ...                 User,
    ...                 (UserProperty.objects, Property.objects, 'properties'),
    ...                 (UserAlias.objects   , Alias.objects   , 'aliases'   ))
    >>> users[0].properties
    [<a list containing Property objects>]
    >>> users[0].aliases
    [<a list containing Alias objects>]
    """
    
    # store the list of object IDs locally for performance
    object_ids = [obj._get_pk_val() for obj in objects]
    
    cached_objects = {}
    
    for (m2m_set, rel_set, attr) in related:
        
        # name of ForeignKey field referring to main objects
        leftfield  = [f.name for f in m2m_set.model._meta.fields
                      if f.rel and f.rel.to._meta is model._meta][0]

        # name of ForeignKey field referring to related objects
        rightfield = [f.name for f in m2m_set.model._meta.fields
                      if f.rel and f.rel.to._meta is rel_set.model._meta][0]

        # query all m2m rows which refer to one of the main objects
        # already in memory
        criteria = {'%s__in' % leftfield: object_ids}
        cached_objects[attr] = {
            'field'  : rightfield,
            'objects': m2m_set.filter(**criteria).select_related(True)}
        
    for obj in objects:
        # store retrieved related objects into corresponding main
        # objects
        for attr, cachedict in cached_objects.iteritems():
            rightfield = cachedict['field']
            event_objects = [getattr(cached_obj, rightfield)
                             for cached_obj in cached_objects[attr]['objects']
                             if getattr(cached_obj, leftfield) == obj]
            setattr(obj, attr, event_objects)
            setattr(obj, attr+'_ids', [eo._get_pk_val() for eo in event_objects])