The Django 2.1 ORM is quite capable. Though it does not cover every conceivable use case, it handles nearly all simple queries and many more complex queries as well. Still, there are times when I wish the ORM had some capability that it doesn’t. That’s happened a couple times this year as I’ve been learning Django. This post tells about a new QuerySet method we developed that’s been useful to us.
The Django ORM allows you to attach data to objects returned from a QuerySet using annotate() and aggregate(). These methods often do exactly what you need. However, there are cases that the Django ORM doesn’t handle. For example, what if I want to attach data that isn’t part of the model or a related model to an object? Perhaps the data is queried from another database or a web API. Maybe the data I need is in an unrelated model and for whatever reason I’m not at liberty to add a relationship. Or perhaps attaching the data via annotate() is theoretically possible but is cumbersome and would be much easier to do in Python code. Another case is attaching data obtained from a raw SQL query: Django has no way to join the result of a raw SQL query to a QuerySet.
From what I can tell, the standard technique in these cases is to use values() to make the QuerySet return a dictionary and then use Python code to add to that dictionary. This can work fine, but I may not want to convert the returned model instances into a dictionary. One example of this is when working with the Django REST Framework, where my serializer expects object instances. Another case is when I want to later call methods on the returned objects, which won’t work once they’ve been converted to dictionaries.
We kept running into situations like this, so I decided to look for a clear and easy to use solution. The result of my search was AugmentableQuerySet. This subclass of QuerySet has a single new public method augment() available. Let me give a trivial example of how to use it.
[code language="py"]
class Product(models.Model):
name = models.CharField(max_length=255)
description = models.TextField()
model_number = IntegerField()
base_price = models.DecimalField(max_digits=8, decimal_places=2)
tax_class = CharField(max_length=1)
objects = AugmentableQuerySet().as_manager()
def format_name_and_model(product):
return f'{product.name} model {product.model_number}'
products = Product.objects.augment(name_and_model=format_name_and_model)
for prod in products:
print(prod.name_and_model)
[/code]
This silly task could be accomplished with annotate(), but it gives an idea of how augment() works: it takes one or more keyword arguments, where the argument name is the field name to be added to the object and the argument value is a function that is called on the object to calculate the value for the new field. It can be anything that is callable: a function, a lambda or a class with a __call__ method. The callable must take a single argument, which is the object (or dictionary) returned by the QuerySet.
Here is a more realistic example that shows you can attach multiple fields to the object.
[code language="py"]
def WarrantyLookup(user):
def lookup(obj):
params = {'country': user.address.country,
'model_number': getattr(obj, 'model_number')}
r = requests.get(WARRANTY_URL, param=params)
return r.json()['period']
return lookup
def product_price(product, user, tax_table):
return complex_calculation_of_tax(product.base_price,
product.tax_class,
user.address.zip_code,
tax_table)
products = Product.objects
.augment(warranty_period=WarrantyLookup(request.user),
price=lambda p: product_price(p, request.user, tax_table))
for prod in products:
print(prod.name, prod.warranty_period, prod.price)
[/code]
The WarrantyLookup is in upper camel case to fit the Django convention for query expressions. (For example, Sum or ExpressionWrapper.) Neither warrant_period nor price are part of the model, but the Product instances returned by the augmented QuerySet have both of those fields.
Below is the implementation of AugmentableQuerySet. It works with QuerySets that return model object instances or dictionaries. It could be modified to handle tuples returned by values_list(). Note that if you call the values() method on the QuerySet to generate dictionaries, the call to augment() must come after it. This is because values() replaces the QuerySet iterator that augment() has wrapped, so the call to augment() will have no effect.
[code language="py"]
class AugmentableQuerySet(models.QuerySet):
@staticmethod
def _make_iterable(iterable, expressions):
class AugmentIterable(iterable):
def __iter__(self):
for row in super().__iter__():
for field, func in expressions.items():
if isinstance(row, dict):
row[field] = func(row)
else:
setattr(row, field, func(row))
yield row
return AugmentIterable
def augment(self, **expressions):
clone = self._chain()
clone._fields = tuple(expressions)
clone._iterable_class = self._make_iterable(self._iterable_class, expressions)
return clone
objects = AugmentableQuerySet().as_manager()[/code]
augment() adds a new field to all objects returned by a queryset, the value of which is the result of applying a function to the object. The implementation is fairly simple, though you don’t need to understand the code above to use augment(). One potential issue is that _iterable_class is not a documented part of Django 2.1, so a new version of Django might break our code. We have some simple tests around the AugmentableQuerySet class that should catch any problems.
We created AugmentableQuerySet to make our Django code shorter and easier to understand. It has been working well for us so far. We may improve the implementation as we gain more experience with it, or perhaps we’ll find a better way to accomplish what we want. For now, though, this has been a successful addition to our Django toolbox.