stackdump/python/packages/sqlobject/util/csvexport.py

"""
Exports a SQLObject class (possibly annotated) to a CSV file.
"""
import os
import csv
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO
import sqlobject

__all__ = ['export_csv', 'export_csv_zip']

def export_csv(soClass, select=None, writer=None, connection=None,
               orderBy=None):
    """
    Export the SQLObject class ``soClass`` to a CSV file.

    ``soClass`` can also be a SelectResult object, as returned by
    ``.select()``.  If it is a class, all objects will be retrieved,
    ordered by ``orderBy`` if given, or the ``.csvOrderBy`` attribute
    if present (but csvOrderBy will only be applied when no select
    result is given).

    You can also pass in select results (or simply a list of
    instances) in ``select`` -- if you have a list of objects (not a
    SelectResult instance, as produced by ``.select()``) then you must
    pass it in with ``select`` and pass the class in as the first
    argument.

    ``writer`` is a ``csv.writer()`` object, or a file-like object.
    If not given, the string of the file will be returned.

    Uses ``connection`` as the data source, if given, otherwise the
    default connection.

    Columns can be annotated with ``.csvTitle`` attributes, which will
    form the attributes of the columns, or 'title' (secondarily), or
    if nothing then the column attribute name.

    If a column has a ``.noCSV`` attribute which is true, then the
    column will be suppressed.

    Additionally a class can have an ``.extraCSVColumns`` attribute,
    which should be a list of strings/tuples.  If a tuple, it should
    be like ``(attribute, title)``, otherwise it is the attribute,
    which will also be the title.  These will be appended to the end
    of the CSV file; the attribute will be retrieved from instances.

    Also a ``.csvColumnOrder`` attribute can be on the class, which is
    the string names of attributes in the order they should be
    presented.
    """

    return_fileobj = None
    if not writer:
        return_fileobj = StringIO()
        writer = csv.writer(return_fileobj)
    elif not hasattr(writer, 'writerow'):
        writer = csv.writer(writer)

    if isinstance(soClass, sqlobject.SQLObject.SelectResultsClass):
        assert select is None, (
            "You cannot pass in a select argument (%r) and a SelectResult argument (%r) for soClass"
            % (select, soClass))
        select = soClass
        soClass = select.sourceClass
    elif select is None:
        select = soClass.select()
        if getattr(soClass, 'csvOrderBy', None):
            select = select.orderBy(soClass.csvOrderBy)

    if orderBy:
        select = select.orderBy(orderBy)
    if connection:
        select = select.connection(connection)

    _actually_export_csv(soClass, select, writer)

    if return_fileobj:
        # They didn't pass any writer or file object in, so we return
        # the string result:
        return return_fileobj.getvalue()

def _actually_export_csv(soClass, select, writer):
    attributes, titles = _find_columns(soClass)
    writer.writerow(titles)
    for soInstance in select:
        row = [getattr(soInstance, attr)
               for attr in attributes]
        writer.writerow(row)

def _find_columns(soClass):
    order = []
    attrs = {}
    for col in soClass.sqlmeta.columnList:
        if getattr(col, 'noCSV', False):
            continue
        order.append(col.name)
        title = col.name
        if hasattr(col, 'csvTitle'):
            title = col.csvTitle
        elif getattr(col, 'title', None) is not None:
            title = col.title
        attrs[col.name] = title

    for attrDesc in getattr(soClass, 'extraCSVColumns', []):
        if isinstance(attrDesc, (list, tuple)):
            attr, title = attrDesc
        else:
            attr = title = attrDesc
        order.append(attr)
        attrs[attr] = title

    if hasattr(soClass, 'csvColumnOrder'):
        oldOrder = order
        order = soClass.csvColumnOrder
        for attr in order:
            if attr not in oldOrder:
                raise KeyError(
                    "Attribute %r in csvColumnOrder (on class %r) does not exist as a column or in .extraCSVColumns (I have: %r)"
                    % (attr, soClass, oldOrder))
            oldOrder.remove(attr)
        order.extend(oldOrder)

    titles = [attrs[attr] for attr in order]
    return order, titles


def export_csv_zip(soClasses, file=None, zip=None, filename_prefix='',
                   connection=None):
    """
    Export several SQLObject classes into a .zip file.  Each
    item in the ``soClasses`` list may be a SQLObject class,
    select result, or ``(soClass, select)`` tuple.

    Each file in the zip will be named after the class name (with
    ``.csv`` appended), or using the filename in the ``.csvFilename``
    attribute.

    If ``file`` is given, the zip will be written to that.  ``file``
    may be a string (a filename) or a file-like object.  If not given,
    a string will be returnd.

    If ``zip`` is given, then the files will be written to that zip
    file.

    All filenames will be prefixed with ``filename_prefix`` (which may
    be a directory name, for instance).
    """
    import zipfile
    close_file_when_finished = False
    close_zip_when_finished = True
    return_when_finished = False
    if file:
        if isinstance(file, basestring):
            close_when_finished = True
            file = open(file, 'wb')
    elif zip:
        close_zip_when_finished = False
    else:
        return_when_finished = True
        file = StringIO()

    if not zip:
        zip = zipfile.ZipFile(file, mode='w')

    try:
        _actually_export_classes(soClasses, zip, filename_prefix,
                                 connection)
    finally:
        if close_zip_when_finished:
            zip.close()
        if close_file_when_finished:
            file.close()

    if return_when_finished:
        return file.getvalue()

def _actually_export_classes(soClasses, zip, filename_prefix,
                             connection):
    for classDesc in soClasses:
        if isinstance(classDesc, (tuple, list)):
            soClass, select = classDesc
        elif isinstance(classDesc, sqlobject.SQLObject.SelectResultsClass):
            select = classDesc
            soClass = select.sourceClass
        else:
            soClass = classDesc
            select = None
        filename = getattr(soClass, 'csvFilename', soClass.__name__)
        if not os.path.splitext(filename)[1]:
            filename += '.csv'
        filename = filename_prefix + filename
        zip.writestr(filename,
                     export_csv(soClass, select, connection=connection))