"""

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))