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    so_class: the SQLObject class to actually stor with 
0039    
0040    """
0041    def __init__(self, so_class):
0042        from sqlobject.styles import getStyle
0043        from sqlobject.dbconnection import Transaction
0044
0045        self.style = getStyle(so_class)
0046        self.so_class = so_class
0047
0048        ## FIXME
0049        ## pg deadlocks in circular fixtures (w/foreign keys)
0050        #if not issubclass(self.so_class._connection.__class__, Transaction):
0051        #    self.so_class._connection = self.so_class._connection.transaction()
0052
0053    def clean(self):
0054        self.so_class.clearTable()
0055
0056    def close(self):
0057        ## FIXME
0058        ## pg deadlocks in circular fixtures (w/foreign keys)
0059        #self.so_class._connection.commit()
0060        pass
0061
0062    def save(self, **kw):
0063        k = [self.style.dbColumnToPythonAttr(k) for k in kw.keys()]
0064        v = kw.values()
0065        self.so_class(**dict(zip(k,v)))
0066
0067class FileStor(Stor):
0068    """abstract Stor object that involves storing a file.
0069    
0070    """
0071    def __init__(self, file):
0072        if type(file) is type(""):
0073            self.file = open(file, 'w')
0074        else:
0075            self.file = file
0076
0077    def _new_file(self, fname):
0078        return open(fname, 'w')
0079
0080    def clean(self):
0081        # is there a better way ?  truncate() no worky
0082        fname = self.file.name
0083        self.file.close()
0084        self.file = self._new_file(fname)
0085
0086    def close(self):
0087        self.file.close()
0088
0089class CsvStor(csv.DictWriter, FileStor):
0090    """CSV storage medium.
0091    
0092    """
0093    def __init__(self, file, fields, addfields=False, encoding=None, clean=False):
0094        """needs a filename or file-like object and list of fields.
0095        
0096        if encoding is set, all values will be encoded with this type
0097        before values are saved
0098        
0099        """
0100        FileStor.__init__(self, file)
0101        self.encoding = encoding
0102        self.addfields = addfields
0103        self.fields = fields
0104        if clean:
0105            FileStor.clean(self)
0106        self._init_csv()
0107
0108    def _init_csv(self):
0109        csv.DictWriter.__init__(self, self.file, self.fields)
0110
0111        if self.addfields:
0112            self.writerow(dict([(f,f) for f in self.fields]))
0113
0114    def clean(self):
0115        FileStor.clean(self)
0116        self._init_csv()
0117
0118    def save(self, **kwds):
0119        if self.encoding:
0120            kwds = dict( [(k.encode(self.encoding), v.encode(self.encoding))                                                           for k,v in kwds.items()] )
0122        self.writerow(kwds)
0123
0124class XmlStor(FileStor):
0125    """XML file storage medium.
0126    
0127    note that this is not meant to get crazy, but just
0128    meant to make an XML representation of *tabular* data ...
0129    in other words the structure needs to be basically: 
0130    
0131        <rows><row><col>...</col></row></rows>
0132    
0133    """
0134    xml = None
0135    item = 'item'
0136    root = 'dataset'
0137    encoding = 'utf-8'
0138    _root_elem = None
0139    _started = False
0140
0141    def __init__(self, file):
0142        if XMLWriter is None:
0143            # force an ImportError
0144            import elementtree.SimpleXMLWriter.XMLWriter
0145
0146        FileStor.__init__(self, file)
0147        self.xml = XMLWriter(self.file, self.encoding)
0148        self._started = False
0149        self._root_elem = None
0150
0151    def close(self):
0152        self.xml.close(self._root_elem)
0153        FileStor.close(self)
0154
0155    def clean(self):
0156        FileStor.clean(self)
0157        # hmm, still not sure about this file-cleaning approach
0158        self.xml = XMLWriter(self.file, self.encoding)
0159        self._started = False
0160        self._root_elem = None
0161
0162    def save(self, **kw):
0163        if not self._started:
0164            self._root_elem = self.xml.start(self.root)
0165            self._started = True
0166
0167        self.start_save()
0168        for k,v in kw.items():
0169            self.xml.element(k, str(v).encode(self.encoding))
0170        self.end_save()
0171
0172    def start_save(self):
0173        self.xml.start(self.item)
0174
0175    def end_save(self):
0176        self.xml.end(self.item)