From 8c715351232e9a2112dc56e0b2a4a96ece75fbd7 Mon Sep 17 00:00:00 2001
From: "brian m. carlson" <sandals@crustytoothpaste.net>
Date: Sun, 11 Jan 2015 17:01:39 +0000
Subject: [PATCH] Add the ability to register a file format backend.

Signed-off-by: brian m. carlson <sandals@crustytoothpaste.net>
---
 lib/newfol/filemanip.py | 22 +++++++++++++---------
 test/testfilemanip.py   | 22 ++++++++++++++++++++++
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/lib/newfol/filemanip.py b/lib/newfol/filemanip.py
index 0fde28d..6e23f93 100644
--- a/lib/newfol/filemanip.py
+++ b/lib/newfol/filemanip.py
@@ -546,23 +546,27 @@ class RawFile(FileFormat):
 class FileStorage:
     """A file (or file-like object) in a certain format."""
     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):
         txnformat = self._canonicalize_transaction_types(txnformat)
         self._txn = self._make_transaction_store(txnformat, options)
-        backends = {
-            "csv": CSVFile,
-            "qddb": QDDBFile,
-            "pickle": PickleFile,
-            "json": JSONFile,
-            "yaml": YAMLFile,
-            "raw": RawFile
-        }
         try:
-            self._backend = backends[fmt](filename, self._txn, options)
+            self._backend = self.BACKENDS[fmt](filename, self._txn, options)
         except KeyError:
             raise ValueError("{0}: not a supported backend".format(fmt))
 
+    @classmethod
+    def register_backend(klass, name, obj):
+        klass.BACKENDS[name] = obj
+
     @staticmethod
     def _canonicalize_transaction_types(types):
         if types is None:
diff --git a/test/testfilemanip.py b/test/testfilemanip.py
index dce504e..091f047 100755
--- a/test/testfilemanip.py
+++ b/test/testfilemanip.py
@@ -241,5 +241,27 @@ class SHA256TransactionTest(unittest.TestCase):
                          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__':
     unittest.main()