Django JSON Helper

json_help.py

"""
Version 0.3.8 - Ed Menendez - ed@digitalhaiku.com

Add to any Django Model definition and you will have a instance.as_json() function).
For example >> class Customer(models.Model, JSONHelper):

Latest at https://gist.github.com/edmenendez/06bcd4ff56a70d2de3fd
"""
import copy
import json

from django.core import serializers
from django.core.exceptions import ObjectDoesNotExist
from django.db.models import fields
from django.utils.deconstruct import deconstructible


@deconstructible
class JSONHelper(object):
    """
    Adds a as_json method to the model.
    """
    def json_cleanup(self, json_dict):
        """
        Override this to take the json_dict and so some additional processing.
        """
        return json_dict

    def as_json(
        self, depth=1, as_python_obj=False, seen_models=None, skip_fields=[],
        include_fields=[],
    ):
        """
        If as_python_obj is True it will not return a json string but a regular
        Python object instead.

        skip_fields should be a list of fields to exclude. include_fields are
        fields to include. If include_fields is blank, it assumes everything is
        included.

        skip_fields currently only works with accessors and FKs but not regular
        fields.
        """
        # House keeping to prevent recursion. If it's seen within the tree, then
        # it skips it.
        seen_models = seen_models and copy.deepcopy(seen_models) or []

        try:
            model = self._meta.model
        except AttributeError:
            # Django 1.5 does this instead
            model = self._meta.concrete_model

        kwargs = {}
        if include_fields:
            kwargs['fields'] = include_fields
        json_obj = serializers.serialize('json', [self, ], **kwargs)

        # Need to add a few extra fields, so lets make it a python object
        json_obj = json.loads(json_obj)[0]

        # Keep track of seen models to prevent endless loops.
        if json_obj['model'] in seen_models:
            #print 'skipping', json_obj['model']
            return json.dumps(['truncated', ]).strip('[]')
        seen_models.append(json_obj['model'])

        # Clean up None fields to empty string
        #for key in json_obj['fields']:
        #    if json_obj['fields'][key] == None:
        #        print key, json_obj['fields'][key]
        #        json_obj['fields'][key] = ''

        # Format date fields.
        for field in model._meta.get_all_field_names():
            try:
                if type(self._meta.get_field(field)) == fields.DateField:
                    json_obj['fields'][field] = getattr(
                        self, field
                    ) and getattr(self, field).strftime('%m/%d/%Y')
            except fields.FieldDoesNotExist:
                pass

        # Add special ForeignKey fields
        for field in set(
            [f.attname for f in model._meta.fields]
        ) - set(model._meta.get_all_field_names()):
            # TODO: This can be optimized with sets above
            if field in skip_fields:
                continue
            if include_fields and field not in include_fields:
                continue

            # Add both the ID and the text in case we're using select2 AJAX lookup
            json_obj['fields'][field] = getattr(self, field)
            json_obj['fields'][field.split('_id')[0]] = unicode(
                getattr(self, unicode(field.split('_id')[0]))
            )

            if depth < 3:
                _obj_key = field.split('_id')[0] + '_obj'
                json_obj['fields'][_obj_key] = 'fail'
                try:
                    json_obj['fields'][_obj_key] = getattr(
                        self, field.split('_id')[0]
                    ).as_json(
                        depth + 1, as_python_obj=True,
                        seen_models=seen_models, skip_fields=skip_fields,
                    )
                    if json_obj['fields'][_obj_key] == '"truncated"':
                        #print 111, depth, json_obj['fields'][field.split('_id')[0] + '_obj']
                        del json_obj['fields'][_obj_key]
                    else:
                        json_obj['fields'][_obj_key] = json.loads(
                            json_obj['fields'][_obj_key]
                        )
                except AttributeError:
                    pass

        # Get all _set fields.
        for related in model._meta.get_all_related_objects():
            accessor_name = related.get_accessor_name()

            # TODO: This can be optimized with sets above
            if accessor_name in skip_fields:
                continue
            if include_fields and accessor_name not in include_fields:
                continue

            try:
                accessor = getattr(self, accessor_name)
            except ObjectDoesNotExist:
                continue

            # Should this be hardcoded at 3? - Ed
            if depth < 3:
                json_obj[accessor_name] = []
                is_not_supported = False
                for instance in accessor.all():
                    try:
                        instance_json = instance.as_json(
                            depth+1, as_python_obj=True,
                            seen_models=seen_models, skip_fields=skip_fields,
                        )
                    except AttributeError:
                        # Do we need to support Django 1.5 here? See above.
                        instance_json = json.dumps(
                            u'%s does not derive from JSONHelper!' %
                            (instance._meta.model._meta.object_name)
                        )
                        is_not_supported = True
                    json_obj[accessor_name].append(json.loads(instance_json))
                    if json_obj[accessor_name][-1] == 'truncated':
                        del json_obj[accessor_name][-1]
                    if is_not_supported:
                        break  # Just quit now
            else:
                json_obj[accessor_name] = json.loads(
                    serializers.serialize('json', accessor.all())
                )

        json_obj['text'] = unicode(self)

        # Hook this here in case some additional manipulation needs to be done.
        json_obj = self.json_cleanup(json_obj)

        if not as_python_obj:
            json_obj = json.dumps([json_obj, ]).strip('[]')

        return json_obj