added XmlStorage class

[sfa.git] / geni / util / storage.py

diff --git a/geni/util/storage.py b/geni/util/storage.py
index ec220b6..bd52351 100644 (file)
--- a/geni/util/storage.py
+++ b/geni/util/storage.py
@@ -1,25 +1,69 @@
+import os
+from geni.util.rspec import RecordSpec
 
 class SimpleStorage(dict):
-
+    """
+    Handles storing and loading python dictionaries. The storage file created
+    is a string representation of the native python dictionary.
+    """
     db_filename = None
-    types = ['dict', 'tabbed', 'text', 'shell']
-
-    def __init__(self, db_filename, db = {}, type = 'dict'):
+    type = 'dict'
+    
+    def __init__(self, db_filename, db = {}):
 
-        if type not in self.types:
-            raise Exception, "Invalid type %s, must be in %s" % (type, self.types)
-        self.type = type
         dict.__init__(self, db)
         self.db_filename = db_filename
     
     def load(self):
-        db_file = open(self.db_filename, 'r')
-        dict.__init__(self, eval(db_file.read()))    
+        if os.path.exists(self.db_filename) and os.path.isfile(self.db_filename):
+            db_file = open(self.db_filename, 'r')
+            dict.__init__(self, eval(db_file.read()))
+        elif os.path.exists(self.db_filename) and not os.path.isfile(self.db_filename):
+            raise IOError, '%s exists but is not a file. please remove it and try again' \
+                           % self.db_filename
+        else:
+            self.write()
+            self.load()
  
     def write(self):
         db_file = open(self.db_filename, 'w')  
-        db_file.write(str(self)
+        db_file.write(str(self))
         db_file.close()
     
     def sync(self):
         self.write()
+
+class XmlStorage(SimpleStorage):
+    """
+    Handles storing and loading python dictionaries. The storage file created
+    is a xml representation of the python dictionary.
+    """ 
+    db_filename = None
+    type = 'xml'
+
+    def load(self):
+        """
+        Parse an xml file and store it as a dict
+        """ 
+        data = RecordSpec()
+        if os.path.exists(self.db_filename) and os.path.isfile(self.db_filename):
+            data.parseFile(self.db_filename)
+            dict.__init__(self, data.toDict())
+        elif os.path.exists(self.db_filename) and not os.path.isfile(self.db_filename):
+            raise IOError, '%s exists but is not a file. please remove it and try again' \
+                           % self.db_filename
+        else:
+            self.write()
+            self.load()
+
+    def write(self):
+        data = RecordSpec()
+        data.parseDict(self)
+        db_file = open(self.db_filename, 'w')
+        db_file.write(data.toprettyxml())
+        db_file.close()
+
+    def sync(self):
+        self.write()
+
+