Writing status uses dump

[moe.git] / t / moe / status.py

diff --git a/t/moe/status.py b/t/moe/status.py
index 936659e6408a907650d8aea60887e6ec5ded5370..5a5adf22c5b4851b572ff955ca08a0814eacb903 100644 (file)
--- a/t/moe/status.py
+++ b/t/moe/status.py
@@ -26,6 +26,9 @@ class MoeStatus:
     def __setitem__(self, k, v):
        self.stat[k] = v
 
+    def keys(self):
+       return self.stat.keys()
+
     def get_list(self, k):
        m = self.stat
        if not m.has_key(k):
@@ -34,25 +37,37 @@ class MoeStatus:
            m[k] = [m[k]]
        return m[k]
 
-    def write(self, file=None, name=None):
-       if file is None:
-           if name is not None:
-               file = open(name, "w")
-           else:
-               file = sys.stdout
-       self.write_nested(file, 0)
-
-    def write_nested(self, file, indent):
+    def dump(self, prefix=""):
+       """
+       Dump status in metafile format.
+       Return a list of lines, `prefix` is indent prefix.
+       """
+       l = []
         for k,v in self.stat.items():
            if type(v) == types.ListType: vals = v
            else: vals = [v]
            for w in vals:
                if isinstance(w, MoeStatus):
-                   file.write("\t" * indent + k + "(\n")
-                   w.write_nested(file, indent+1);
-                   file.write("\t" * indent + ")\n")
+                   l.append(prefix + k + "(\n")
+                   l.extend(self.str_lines(prefix+"  "))
+                   l.append(prefix + ")\n")
                else:
-                   file.write("\t" * indent + k + ":" + str(w) + "\n")
+                   l.append(prefix + k + ":" + str(w) + "\n")
+       return l
+       
+    def write(self, f=None, name=None):
+       """
+       Write status (as a metafile) to `f` or file `name` or `stdout`
+       """
+       if not f and name is not None:
+           with open(name, "w") as f:
+               for l in self.dump():
+                   f.write(l+"\n")
+       else:
+           if not f: 
+               f = sys.stdout
+           for l in self.dump():
+               f.write(l+"\n")
 
     def read(self, file=None, name=None):
        if file is None: