0001# Copyright (C) 2006 Kumar McMillan
0002
0003# This library is free software; you can redistribute it and/or
0004# modify it under the terms of the GNU Lesser General Public
0005# License as published by the Free Software Foundation; either
0006# version 2.1 of the License, or (at your option) any later version.
0007
0008"""Fixture storage media.
0009
0010You most likely will only need the wrapper classes in `fixtures`.
0011
0012"""
0013
0014import csv
0015import sys
0016try:
0017    from elementtree.SimpleXMLWriter import XMLWriter
0018except ImportError:
0019    XMLWriter = None
0020from exceptions import *
0021from config import debug
0022
0023class Stor(object):
0024    """interface for a storage medium."""
0025    def clean(self):
0026        """erases any previous contents in storage."""
0027        raise NotImplementedError
0028    def close(self):
0029        """perform any close logic, like closing a file."""
0030        raise NotImplementedError
0031    def save(self, **kwds):
0032        """saves fields in kwds."""
0033        raise NotImplementedError
0034
0035class SOStor(Stor):
0036    """SQLObject storage medium.
0037    
0038    Arguments
0039    ---------
0040    
0041    so_class: the SQLObject class to actually stor with 
0042    
0043    Keyword Arguments
0044    -----------------
0045    
0046    autocommit -- when False (the default) everything is performed within \
0047        a transaction until `SOStor.clean` is called.  When True, \
0048        everything is still within a transaction but it is committed \
0049        upon `SOStor.save`.  Autcommit mode seems to be necessary to work \
0050        around instances where postgresql would deadlock if you are working \
0051        with lots of foreign keyed fixtures within transactions
0052    
0053    """
0054    _trans_per_conn = {} # transactions need to be singletons per 
0055                         # dsn, so that we can work with multiple objects 
0056                         # (i.e. foreign keys)
0057                         # and ... we need transactions to work around 
0058                         # issues with sqlite in sqlobject
0059
0060    def __init__(self, so_class, autocommit=False):
0061        from sqlobject.styles import getStyle
0062        from sqlobject.dbconnection import Transaction
0063
0064        self.style = getStyle(so_class)
0065        self.so_class = so_class
0066        self.autocommit = autocommit
0067
0068        cls = self.__class__
0069        conn = self.so_class._connection
0070        if (not hasattr(conn, 'dsn') and
0071                hasattr(conn, 'filename')):
0072            # sqlite ..
0073            connid = conn.filename
0074        else:
0075            # postgres and ..?..
0076            connid = conn.dsn
0077
0078        if not cls._trans_per_conn.has_key(connid):
0079            cls._trans_per_conn[connid] = conn.transaction()
0080        self._trans = cls._trans_per_conn[connid]
0081
0082    def __repr__(self):
0083        return "<class '%s' so='%s' at %s>" %                           (self.__class__.__name__,
0085                         self.so_class.__name__,
0086                         hex(id(self)))
0087
0088    def clean(self):
0089        self.so_class.clearTable()
0090
0091    def close(self):
0092        self._trans.commit()
0093
0094    def save(self, **kw):
0095        k = [self.style.dbColumnToPythonAttr(k) for k in kw.keys()]
0096        v = kw.values()
0097        new_kw = dict(zip(k,v))
0098        new_kw['connection'] = self._trans
0099        self.so_class(**new_kw)
0100        if self.autocommit:
0101            self._trans.commit()
0102            try:
0103                self._trans.begin()
0104            except AssertionError:
0105                # if we are in a bad transaction
0106                # we should just swallow the "transaction is _obsolete" error?
0107                pass
0108
0109class FileStor(Stor):
0110    """abstract Stor object that involves storing a file.
0111    
0112    """
0113    def __init__(self, file):
0114        if type(file) is type(""):
0115            self.file = open(file, 'w')
0116        else:
0117            self.file = file
0118
0119    def _new_file(self, fname):
0120        return open(fname, 'w')
0121
0122    def clean(self):
0123        # is there a better way ?  truncate() no worky
0124        fname = self.file.name
0125        self.file.close()
0126        self.file = self._new_file(fname)
0127
0128    def close(self):
0129        self.file.close()
0130
0131class CsvStor(csv.DictWriter, FileStor):
0132    """CSV storage medium.
0133    
0134    """
0135    def __init__(self, file, fields, addfields=False,
0136                                encoding=None, clean=False, dialect='excel'):
0137        """needs a filename or file-like object and list of fields.
0138        
0139        if encoding is set, all values will be encoded with this type
0140        before values are saved
0141        
0142        """
0143        FileStor.__init__(self, file)
0144        self.encoding = encoding
0145        self.addfields = addfields
0146        self.fields = fields
0147        self.dialect = dialect
0148        if clean:
0149            FileStor.clean(self)
0150        self._init_csv()
0151
0152    def _init_csv(self):
0153        csv.DictWriter.__init__(self, self.file,
0154                                    self.fields, dialect=self.dialect)
0155
0156        if self.addfields:
0157            self.writerow(dict([(f,f) for f in self.fields]))
0158
0159    def clean(self):
0160        FileStor.clean(self)
0161        self._init_csv()
0162
0163    def save(self, **kwds):
0164        if self.encoding:
0165            kwds = dict(
0166                  [(k.encode(self.encoding),
0167                    v.encode(self.encoding)) for k,v in kwds.items()] )
0168        self.writerow(kwds)
0169
0170class XmlStor(FileStor):
0171    """XML file storage medium.
0172    
0173    note that this is not meant to get crazy, but just
0174    meant to make an XML representation of *tabular* data ...
0175    in other words the structure needs to be basically: 
0176    
0177        <rows><row><col>...</col></row></rows>
0178    
0179    """
0180
0181    def __init__(self, file, encoding='utf-8', item='item', root='dataset'):
0182        if XMLWriter is None:
0183            # force an ImportError
0184            import elementtree.SimpleXMLWriter.XMLWriter
0185
0186        FileStor.__init__(self, file)
0187        # note that 'us-ascii' is elementtree default
0188        self.item = item
0189        self.root = root
0190        self.encoding = encoding
0191        self.xml = XMLWriter(self.file, self.encoding)
0192        self._started = False
0193        self._root_elem = None
0194
0195    def close(self):
0196        self.xml.close(self._root_elem)
0197        FileStor.close(self)
0198
0199    def clean(self):
0200        FileStor.clean(self)
0201        # hmm, still not sure about this file-cleaning approach
0202        self.xml = XMLWriter(self.file, self.encoding)
0203        self.xml.declaration()
0204        self._started = False
0205        self._root_elem = None
0206
0207    def save(self, **kw):
0208        if not self._started:
0209            self._root_elem = self.xml.start(self.root)
0210            self._started = True
0211
0212        self.start_save()
0213        for k,v in kw.items():
0214            self.xml.element(k, str(v).encode(self.encoding))
0215        self.end_save()
0216
0217    def start_save(self):
0218        self.xml.start(self.item)
0219
0220    def end_save(self):
0221        self.xml.end(self.item)