0001
0002
0003
0004
0005
0006
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 = {}
0055
0056
0057
0058
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
0073 connid = conn.filename
0074 else:
0075
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
0106
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
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
0184 import elementtree.SimpleXMLWriter.XMLWriter
0185
0186 FileStor.__init__(self, file)
0187
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
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)