|
From: <sub...@co...> - 2005-10-26 22:58:06
|
Author: test Date: 2005-10-26 22:57:59 +0000 (Wed, 26 Oct 2005) New Revision: 1143 Added: SQLObject/trunk/sqlobject/tests/test_csvexport.py SQLObject/trunk/sqlobject/util/csvexport.py Modified: SQLObject/trunk/docs/SQLObject.txt Log: Added routine for doing CSV exports Modified: SQLObject/trunk/docs/SQLObject.txt =================================================================== --- SQLObject/trunk/docs/SQLObject.txt 2005-10-26 22:40:08 UTC (rev 1142) +++ SQLObject/trunk/docs/SQLObject.txt 2005-10-26 22:57:59 UTC (rev 1143) @@ -1668,6 +1668,15 @@ and ALL(). The last 3 are used with comparison operators, like this: "somevalue = ANY(Select(...))". +Utilities +--------- + +Some useful utility functions are included with SQLObject. For more +information see their module docstrings. + +* `sqlobject.util.csvexport + <http://sqlobject.org/module-sqlobject.util.csvexport.html>`_ + SQLBuilder ---------- Added: SQLObject/trunk/sqlobject/tests/test_csvexport.py =================================================================== --- SQLObject/trunk/sqlobject/tests/test_csvexport.py 2005-10-26 22:40:08 UTC (rev 1142) +++ SQLObject/trunk/sqlobject/tests/test_csvexport.py 2005-10-26 22:57:59 UTC (rev 1143) @@ -0,0 +1,90 @@ +from sqlobject import * +from dbtest import * +import csv +from StringIO import StringIO +from sqlobject.util.csvexport import export_csv, export_csv_zip + +def assert_export(result, *args, **kw): + f = StringIO() + kw['writer'] = f + export_csv(*args, **kw) + s = f.getvalue().replace('\r\n', '\n') + if result.strip() != s.strip(): + print '**Expected:' + print result + print '**Got:' + print s + assert result.strip() == s.strip() + +class Simple(SQLObject): + + name = StringCol() + address = StringCol() + address.csvTitle = 'Street Address' + hidden = StringCol() + hidden.noCSV = True + +class Complex(SQLObject): + + fname = StringCol() + lname = StringCol() + age = IntCol() + extraCSVColumns = [('name', 'Full Name'), 'initials'] + # initials should end up at the end then: + csvColumnOrder = ['name', 'fname', 'lname', 'age'] + def _get_name(self): + return self.fname + ' ' + self.lname + def _get_initials(self): + return self.fname[0] + self.lname[0] + +def test_simple(): + setupClass(Simple) + Simple(name='Bob', address='432W', hidden='boo') + Simple(name='Joe', address='123W', hidden='arg') + assert_export("""\ +name,Street Address +Bob,432W +Joe,123W +""", Simple, orderBy='name') + assert_export("""\ +name,Street Address +Joe,123W +Bob,432W +""", Simple, orderBy='address') + assert_export("""\ +name,Street Address +Joe,123W +""", Simple.selectBy(name='Joe')) + +def test_complex(): + setupClass(Complex) + Complex(fname='John', lname='Doe', age=40) + Complex(fname='Bob', lname='Dylan', age=60) + Complex(fname='Harriet', lname='Tubman', age=160) + assert_export("""\ +Full Name,fname,lname,age,initials +John Doe,John,Doe,40,JD +Bob Dylan,Bob,Dylan,60,BD +Harriet Tubman,Harriet,Tubman,160,HT +""", Complex, orderBy='lname') + assert_export("""\ +Full Name,fname,lname,age,initials +Bob Dylan,Bob,Dylan,60,BD +John Doe,John,Doe,40,JD +""", Complex.select(Complex.q.lname.startswith('D'), orderBy='fname')) + +def test_zip(): + # Just exercise tests, doesn't actually test results + setupClass(Simple) + Simple(name='Bob', address='432W', hidden='boo') + Simple(name='Joe', address='123W', hidden='arg') + + setupClass(Complex) + Complex(fname='John', lname='Doe', age=40) + Complex(fname='Bob', lname='Dylan', age=60) + Complex(fname='Harriet', lname='Tubman', age=160) + s = export_csv_zip([Simple, Complex]) + assert isinstance(s, str) and s + s = export_csv_zip([Simple.selectBy(name='Bob'), + (Complex, list(Complex.selectBy(fname='John')))]) + Added: SQLObject/trunk/sqlobject/util/csvexport.py =================================================================== --- SQLObject/trunk/sqlobject/util/csvexport.py 2005-10-26 22:40:08 UTC (rev 1142) +++ SQLObject/trunk/sqlobject/util/csvexport.py 2005-10-26 22:57:59 UTC (rev 1143) @@ -0,0 +1,196 @@ +""" +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 not select, ( + "You cannot pass in a select argument (%r) and a SelectResult argument (%r) for soClass" + % (select, soClass)) + select = soClass + soClass = select.sourceClass + elif not select: + 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, (str, unicode)): + 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)) |
