Add the ability to register a file format backend.

Signed-off-by: brian m. carlson <sandals@crustytoothpaste.net>
This commit is contained in:
brian m. carlson 2015-01-11 17:01:39 +00:00
parent 72c30291fa
commit 8c71535123
No known key found for this signature in database
GPG key ID: BF535D811F52F68B
2 changed files with 35 additions and 9 deletions

View file

@ -546,23 +546,27 @@ class RawFile(FileFormat):
class FileStorage: class FileStorage:
"""A file (or file-like object) in a certain format.""" """A file (or file-like object) in a certain format."""
DEFAULT_TRANSACTION_TYPES = ('hash',) DEFAULT_TRANSACTION_TYPES = ('hash',)
BACKENDS = {
"csv": CSVFile,
"qddb": QDDBFile,
"pickle": PickleFile,
"json": JSONFile,
"yaml": YAMLFile,
"raw": RawFile
}
def __init__(self, fmt, filename, txnformat=None, options=None): def __init__(self, fmt, filename, txnformat=None, options=None):
txnformat = self._canonicalize_transaction_types(txnformat) txnformat = self._canonicalize_transaction_types(txnformat)
self._txn = self._make_transaction_store(txnformat, options) self._txn = self._make_transaction_store(txnformat, options)
backends = {
"csv": CSVFile,
"qddb": QDDBFile,
"pickle": PickleFile,
"json": JSONFile,
"yaml": YAMLFile,
"raw": RawFile
}
try: try:
self._backend = backends[fmt](filename, self._txn, options) self._backend = self.BACKENDS[fmt](filename, self._txn, options)
except KeyError: except KeyError:
raise ValueError("{0}: not a supported backend".format(fmt)) raise ValueError("{0}: not a supported backend".format(fmt))
@classmethod
def register_backend(klass, name, obj):
klass.BACKENDS[name] = obj
@staticmethod @staticmethod
def _canonicalize_transaction_types(types): def _canonicalize_transaction_types(types):
if types is None: if types is None:

View file

@ -241,5 +241,27 @@ class SHA256TransactionTest(unittest.TestCase):
None) None)
class ExampleFile(newfol.filemanip.FileFormat):
def __init__(self, filename, txn, options):
self.inited = True
def store(self, records):
self.stored = True
pass
def load(self):
return [Record(['a', 'b'])]
class PluggableBackendsTest(unittest.TestCase):
def test_adding_backend(self):
FileStorage.register_backend('example', ExampleFile)
self.assertEqual(FileStorage.BACKENDS['example'], ExampleFile)
fs = FileStorage('example', 'filename')
recs = fs.load()
self.assertEqual(len(recs), 1)
self.assertEqual(recs[0].fields, ['a', 'b'])
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()