pk_lookup.py

'''
Shorthand decorator for the common case of looking up ORM objects from
named URL parameters.

For each named URL parameter passed to the view of the form
``db_table_id``, look up the ORM object in the corresponding model with
the given primary key, returning a 404 in case it is not found.

For example:

    # URL:
      url(r'^/products/(?P<product_id>\d+)/$', views.product_detail,
          name='product_detail'),

    # View:
    @pk_lookup
    def product_detail(request, product):
        # ...
        return direct_to_template(request, 'product_detail.html',
                                  {'product': product})

You can optionally supply a list of filters to limit the queryset
selected from (which is by default the entire model). This is done by
passing parameters to the decorator, using the same syntax as the queryset
method ``filter``.

    @pk_lookup(hidden=False)
    def product_detail(request, product):
        # ...

'''

from functools import wraps

from django.db.models import get_models
from django.shortcuts import get_object_or_404

ID_FIELDS = {}
for model in get_models():
    ID_FIELDS['%s_id' % model._meta.db_table] = (model._meta.db_table,
                                                 model)

def pk_lookup(*args, **filters):
    assert (not filters and len(args) == 1) or (not args and filters)
    def decorator(view):
        def wrapped(request, **kwargs):
            for arg, val in kwargs.iteritems():
                if arg in ID_FIELDS:
                    table, model = ID_FIELDS[arg]
                    del kwargs[arg]
                    kwargs[str(table)] = get_object_or_404(model, pk=val,
                                                           **filters)
            return view(request, **kwargs)
        return wrapped
    if args:
        view = args[0]
        return wraps(view)(decorator(view))
    return decorator