export_csv_command.py

# A reusable Django Export CSV Command
# https://gist.github.com/bryanchow/e611598bc9391ac854ff427582e5619c
#
# Usage: Import this module in yourapp/management/commands/yourcommand.py
# and subclass ExportCSVCommand like this:
#
# class Command(ExportCSVCommand):
#     def get_queryset(self):
#         return YourModel.objects.all()
#     def handle_exception(self, e, inst):
#         do_something_with(e)

import csv
from django.core.management.base import BaseCommand


class ExportCSVCommand(BaseCommand):

    help = "Generate a CSV export of model data."

    def get_queryset(self):
        raise RuntimeError, "ExportCSVCommand is meant to be subclassed."

    def add_arguments(self, parser):
        parser.add_argument(
            "filepath",
            type = str,
            help = "The path of the CSV file to be written",
        )
        parser.add_argument(
            "fields",
            type = str,
            help = "A comma-separated list of field names to include",
        )
        parser.add_argument(
            "--header",
            action = 'store_false',
            help = "Write a header row containing the field names",
        )

    def handle(self, *args, **options):

        filepath = options.get('filepath')
        fields = options.get('fields')
        header = options.get('header')

        field_names = [x.strip() for x in fields.split(",")]
        queryset = self.get_queryset()

        with open(filepath, 'w') as csvfile:

            writer = csv.writer(csvfile)

            if header:
                header_row = []
                for name in field_names:
                    try:
                        field = getattr(queryset.model, name)
                        if hasattr(field, 'short_description'):
                            header_row.append(field.short_description)
                        else:
                            header_row.append(name)
                    except AttributeError:
                        header_row.append(name)
                writer.writerow(header_row)

            for inst in queryset:
                try:
                    row = []
                    for field in field_names:
                        value = getattr(inst, field)
                        if callable(value):
                            value = value()
                        else:
                            value = "" if value is None else value
                        try:
                            # Python 2
                            row.append(unicode(value).encode("utf-8"))
                        except NameError:
                            # Python 3
                            row.append(value)
                    writer.writerow(row)
                except Exception as e:
                    self.handle_exception(e, inst)
                    continue

    def handle_exception(self, e, inst):
        raise e