Version 0.4: Added action to export selected records

[auf_django_export.git] / auf / django / export / admin.py

# -*- encoding: utf-8 -*-

import csv
from StringIO import StringIO

from odf.opendocument import OpenDocumentSpreadsheet
from odf.style import Style, TextProperties, TableColumnProperties
from odf.text import P
from odf.table import Table, TableRow, TableCell

from django.core.urlresolvers import reverse
from django.contrib import admin
from django.utils.functional import update_wrapper
from django.http import HttpResponse

# Create a notation similar to queryset filter
SEPARATOR = u"__"

# String return if lookup property failed
EMPTY = u""


def txt(msg):
    text = msg
    if not isinstance(msg, unicode):
        try:
            text = unicode (msg, "utf-8")
        except TypeError:
            pass
    return text


class ExportAdmin(admin.ModelAdmin):
    change_list_template = 'admin/export/change_list.html'
    actions = ['export_csv', 'export_ods']

    def changelist_view(self, request, extra_context=None):
        """
        Add export URL for the template
        """
        if extra_context is None:
            extra_context = {}
        extra_context.update({
            'export_csv_url' : reverse("admin:%s_%s_export_csv" % (self.model._meta.app_label, self.model._meta.module_name)),
            'export_ods_url' : reverse("admin:%s_%s_export_ods" % (self.model._meta.app_label, self.model._meta.module_name)),
        })
        return super(ExportAdmin, self).changelist_view(request, extra_context)


    def get_urls(self):
        from django.conf.urls.defaults import patterns, url

        urls = super(ExportAdmin, self).get_urls()
        info = self.model._meta.app_label, self.model._meta.module_name

        urlpatterns = patterns('',
            url(r'^export-csv/$',
                self.admin_site.admin_view(self.export_csv),
                name='%s_%s_export_csv' % info),
            url(r'^export-ods/$',
                self.admin_site.admin_view(self.export_ods),
                name='%s_%s_export_ods' % info),
        )
        return urlpatterns + urls

    def get_object_value(self, obj, header):
        """
        Function which work as a queryset filter access :
        Father -- Kid (One2One father related_name = 'child')
        father__name
        father__child__name
        It tries to resolve the property from header string, through
        linked object if needed. If it can't load any value, empty
        string is return.
        """
        segments = header.split(SEPARATOR)
        segments.reverse()
        prop = segments.pop()
        if len(segments) == 0:
            return getattr(obj, prop, EMPTY)
        else:
            try:
                child = getattr(obj, prop, None)
                return self.get_object_value(child, ".".join(segments))
            except:
                return EMPTY

    def get_default_export_fields(self):
      fields_name = [f.name for f in  self.model._meta.fields]
      for fk in self.model._meta._related_objects_cache.keys():
          fields_name = fields_name + [u"%s%s%s" % (fk.var_name, SEPARATOR, f.name) for f in fk.model._meta.fields]
      return fields_name


    def get_export_fields(self):
        """
        This property is a list of string, which should have queryset filter notation style (with X2 underscore).
        If there is no fields, the models is full introspected, stop at 1 depth level.
        """

        export_fields = getattr(self, 'export_fields', None)

        if export_fields is None:
            export_fields = getattr(self, 'csv_fields', None)

        if export_fields is None:
            export_fields = self.get_default_export_fields()

        return export_fields

    def export_ods(self, request, queryset=None):
        """
        Generate HTTP response as ODS file from export_fields property.
        """
        export_fields = self.get_export_fields()

        if queryset is None:
            queryset = self.queryset(request)

        response = HttpResponse(mimetype='application/vnd.oasis.opendocument.spreadsheet')
        response['Content-Disposition'] = 'attachment; filename=%s.ods' % unicode(self.model._meta.verbose_name_plural)

        doc = OpenDocumentSpreadsheet()
        style = Style(name="Large number", family="table-cell")
        style.addElement(TextProperties(fontfamily="Arial", fontsize="15pt"))
        doc.styles.addElement(style)
        widewidth = Style(name="co1", family="table-column")
        widewidth.addElement(TableColumnProperties(columnwidth="2.8cm", breakbefore="auto"))
        doc.automaticstyles.addElement(widewidth)

        table = Table()
        if len (export_fields) > 0:
            tr = TableRow ()
            table.addElement (tr)
            for item in export_fields:
                tc = TableCell ()
                tr.addElement (tc)
                p = P(stylename = style, text = txt(item))
                tc.addElement (p)

        for o in queryset:
            tr = TableRow ()
            table.addElement (tr)

            for attr in export_fields:
                value = self.get_object_value(o, attr)
                tc = TableCell ()
                tr.addElement (tc)
                p = P (stylename = style, text = txt(value))
                tc.addElement (p)

        doc.spreadsheet.addElement(table)
        buffer = StringIO ()
        doc.write(buffer)

        response.write(buffer.getvalue())
        return response

    def export_csv(self, request, queryset=None):
        """
        Generate HTTP response as CSV file from export_fields property.
        """
        csv_fields = self.get_export_fields()

        if queryset is None:
            queryset = self.queryset(request)

        response = HttpResponse(mimetype='text/csv')
        response['Content-Disposition'] = 'attachment; filename=%s.csv' % unicode(self.model._meta.verbose_name_plural)

        writer = csv.writer(response)

        headers = []
        for attr in csv_fields:
            headers.append(attr.encode('utf-8') if isinstance(attr, unicode) else attr)
        writer.writerow(headers)

        for o in queryset:
            row = []
            for attr in csv_fields:
                value = self.get_object_value(o, attr)
                row.append(value.encode('utf-8') if isinstance(value, unicode) else value)
            writer.writerow(row)

        return response