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 SAStor(Stor):
0110 def __init__(self, mapper):
0111 from sqlalchemy import create_session
0112 self.mapper = mapper
0113 self.session = create_session()
0114
0115 def clean(self):
0116 pass
0117
0118 def close(self):
0119 self.session.flush()
0120
0121 def save(self, **kw):
0122 obj = self.mapper()
0123 [setattr(obj, k,v) for k,v in kw.items()]
0124 self.session.save(obj)
0125
0126class FileStor(Stor):
0127 """abstract Stor object that involves storing a file.
0128
0129 """
0130 def __init__(self, file):
0131 if type(file) is type(""):
0132 self.file = open(file, 'w')
0133 else:
0134 self.file = file
0135
0136 def _new_file(self, fname):
0137 return open(fname, 'w')
0138
0139 def clean(self):
0140
0141 fname = self.file.name
0142 self.file.close()
0143 self.file = self._new_file(fname)
0144
0145 def close(self):
0146 self.file.close()
0147
0148class CsvStor(csv.DictWriter, FileStor):
0149 """CSV storage medium.
0150
0151 """
0152 def __init__(self, file, fields, addfields=False,
0153 encoding=None, clean=False, dialect='excel'):
0154 """needs a filename or file-like object and list of fields.
0155
0156 if encoding is set, all values will be encoded with this type
0157 before values are saved
0158
0159 """
0160 FileStor.__init__(self, file)
0161 self.encoding = encoding
0162 self.addfields = addfields
0163 self.fields = fields
0164 self.dialect = dialect
0165 if clean:
0166 FileStor.clean(self)
0167 self._init_csv()
0168
0169 def _init_csv(self):
0170 csv.DictWriter.__init__(self, self.file,
0171 self.fields, dialect=self.dialect)
0172
0173 if self.addfields:
0174 self.writerow(dict([(f,f) for f in self.fields]))
0175
0176 def clean(self):
0177 FileStor.clean(self)
0178 self._init_csv()
0179
0180 def save(self, **kwds):
0181 if self.encoding:
0182 kwds = dict(
0183 [(k.encode(self.encoding),
0184 v.encode(self.encoding)) for k,v in kwds.items()] )
0185 self.writerow(kwds)
0186
0187class XmlStor(FileStor):
0188 """XML file storage medium.
0189
0190 note that this is not meant to get crazy, but just
0191 meant to make an XML representation of *tabular* data ...
0192 in other words the structure needs to be basically:
0193
0194 <rows><row><col>...</col></row></rows>
0195
0196 """
0197
0198 def __init__(self, file, encoding='utf-8', item='item', root='dataset'):
0199 if XMLWriter is None:
0200
0201 import elementtree.SimpleXMLWriter.XMLWriter
0202
0203 FileStor.__init__(self, file)
0204
0205 self.item = item
0206 self.root = root
0207 self.encoding = encoding
0208 self.xml = XMLWriter(self.file, self.encoding)
0209 self._started = False
0210 self._root_elem = None
0211
0212 def close(self):
0213 self.xml.close(self._root_elem)
0214 FileStor.close(self)
0215
0216 def clean(self):
0217 FileStor.clean(self)
0218
0219 self.xml = XMLWriter(self.file, self.encoding)
0220 self.xml.declaration()
0221 self._started = False
0222 self._root_elem = None
0223
0224 def save(self, **kw):
0225 if not self._started:
0226 self._root_elem = self.xml.start(self.root)
0227 self._started = True
0228
0229 self.start_save()
0230 for k,v in kw.items():
0231 self.xml.element(k, str(v).encode(self.encoding))
0232 self.end_save()
0233
0234 def start_save(self):
0235 self.xml.start(self.item)
0236
0237 def end_save(self):
0238 self.xml.end(self.item)