t/moe/status.py

   1 import sys
   2 import types
   3 import re
   4
   5 key_pattern = re.compile("\A[A-Za-z0-9_-]+\Z")
   6
   7 class MoeInvalidStatusFile(Exception):
   8     pass
   9
  10 class Status:
  11     """
  12     (One subtree of) Moe status file.
  13     """
  14
  15     def __init__(self):
  16         self.d = {}
  17
  18     def __getitem__(self, k):
  19         return self.d[k]
  20
  21     def __setitem__(self, k, v):
  22         self.d[k] = v
  23
  24     def keys(self):
  25         return self.d.keys()
  26
  27     def update(self, stat2):
  28         """
  29         Updates values of `self` with values of `stat2`, recursively.
  30
  31         Directly references objects (values and subtrees) of `stat2`, so making a deep copy of `stat2`
  32         may be necessary if you intend to modify `stat2` afterwards.
  33         """
  34
  35         for k,v2 in stat2.d.items():
  36             if k not in self.d:
  37                 self[k] = v2
  38             else:
  39                 v = self[k]
  40                 if isinstance(v, Status) != isinstance(v2, Status):
  41                     raise TypeError("Mixing Status and value while updating key %r"%k)
  42                 if isinstance(v, Status):
  43                     v.update(v2)
  44                 else:
  45                     self[k] = v2
  46
  47     def dump(self, prefix=""):
  48         """
  49         Dump Status in status file format.
  50         Returns a list of lines, ``prefix`` is indentation prefix.
  51         """
  52
  53         l = []
  54         for k,v in self.d.items():
  55             if isinstance(v, Status):
  56                 l.append(prefix + k + " (")
  57                 l.extend(self.dump(prefix+"  "))
  58                 l.append(prefix + ")")
  59             else:
  60                 d = str(v).split('\n')
  61                 l.append(prefix + k + ":" + d[0])
  62                 for i in d[1:]:
  63                     l.append(prefix + ' '*len(k) + ':' + i)
  64         return l
  65
  66     def write(self, f=None, name=None):
  67         """
  68         Write Status to File ``f`` or overwrite file ``name`` or write to ``stdout`` (otherwise).
  69         """
  70
  71         if not f and name is not None:
  72             with open(name, "w") as f:
  73                 for l in self.dump():
  74                     f.write(l+"\n")
  75         else:
  76             if not f:
  77                 f = sys.stdout
  78             for l in self.dump():
  79                 f.write(l+"\n")
  80
  81     def read(self, f=None, name=None):
  82         """
  83         Parse Status from File ``f`` or from file ``name`` or from ``stdin`` (otherwise)
  84         Deletes all contents of the Status.
  85         """
  86
  87         self.d = {}
  88         if not f and name is not None:
  89             with open(name, 'r') as f:
  90                 self.do_read(f)
  91         else:
  92             if not f:
  93                 f = sys.stdin
  94             self.do_read(f)
  95
  96     def read_val(self, k, v):
  97         """
  98         Internal: Safely add a new value
  99         """
 100
 101         if not key_pattern.match(k):
 102             raise MoeInvalidStatusFile("Parse error: invalid key %r"%k)
 103         if k in self.d:
 104             raise MoeInvalidStatusFile("Multiple occurences of key %r"%k)
 105         self.d[k]=v
 106
 107     def do_read(self, f):
 108         """
 109         Internal: Parse an open file
 110         """
 111
 112         stk = []
 113         this = self
 114         for x in f.readlines():
 115             x = x.rstrip("\n").lstrip(" \t")
 116             if x=="" or x.startswith("#"):
 117                 pass
 118             else:
 119                 sep = x.find(":")
 120                 if sep >= 0:
 121                     k = x[:sep].rstrip(" \t")
 122                     v = x[sep+1:]
 123                     this.read_val(k, v)
 124                 elif x.endswith("("):
 125                     k = x[:-1].rstrip(" \t")
 126                     new = Status()
 127                     this.read_val(k, new)
 128                     stk.append(this)
 129                     this = new
 130                 elif x == ")":
 131                     if len(stk) == 0:
 132                         raise MoeInvalidStatusFile("Parse error: incorrect nesting")
 133                     else:
 134                         this = stk.pop()
 135                 else:
 136                     raise MoeInvalidStatusFile("Parse error: malformed line")