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 InvalidStatusFile(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 __eq__(self, s):
  25         return self.d == s.d
  26
  27     def keys(self):
  28         return self.d.keys()
  29
  30     def update(self, stat2):
  31         """
  32         Updates values of `self` with values of `stat2`, recursively.
  33
  34         Directly references objects (values and subtrees) of `stat2`, so making a deep copy of `stat2`
  35         may be necessary if you intend to modify `stat2` afterwards.
  36         """
  37
  38         for k,v2 in stat2.d.items():
  39             if k not in self.d:
  40                 self[k] = v2
  41             else:
  42                 v = self[k]
  43                 if isinstance(v, Status) != isinstance(v2, Status):
  44                     raise TypeError("Mixing Status and value while updating key %r"%k)
  45                 if isinstance(v, Status):
  46                     v.update(v2)
  47                 else:
  48                     self[k] = v2
  49
  50     def dump(self, prefix=""):
  51         """
  52         Dump Status in status file format.
  53         Returns a list of lines, ``prefix`` is indentation prefix.
  54         """
  55
  56         l = []
  57         for k,v in self.d.items():
  58             if isinstance(v, Status):
  59                 l.append(prefix + k + " (")
  60                 l.extend(v.dump(prefix+"  "))
  61                 l.append(prefix + ")")
  62             else:
  63                 d = str(v).split('\n')
  64                 l.append(prefix + k + ":" + d[0])
  65                 for i in d[1:]:
  66                     l.append(prefix + ' '*len(k) + ':' + i)
  67         return l
  68
  69     def write(self, f=None, name=None):
  70         """
  71         Write Status to File ``f`` or overwrite file ``name`` or write to ``stdout`` (otherwise).
  72         """
  73
  74         if not f and name is not None:
  75             with open(name, "w") as f:
  76                 for l in self.dump():
  77                     f.write(l+"\n")
  78         else:
  79             if not f:
  80                 f = sys.stdout
  81             for l in self.dump():
  82                 f.write(l+"\n")
  83
  84     def read(self, f=None, name=None, lines=None):
  85         """
  86         Parse Status file
  87         * from File ``f``
  88         * or from file ``name`` opened for reading 8-bit ASCII
  89         * or from ``lines`` (a list/iterator of lines)
  90
  91         Deletes all previous contents of the Status.
  92         """
  93
  94         self.d = {}
  95         if f is not None:
  96             return self.do_read(f.readlines())
  97         if name is not None:
  98             with open(name, 'r') as f:
  99                 return self.do_read(f.readlines())
 100         if lines is not None:
 101             return self.do_read(lines)
 102         raise ValueError('Provide at least one parameter to Status.read()')
 103
 104     def read_val(self, k, v):
 105         """
 106         Internal: Safely add a new value to Status
 107         """
 108
 109         if not key_pattern.match(k):
 110             raise InvalidStatusFile("Parse error: invalid key %r"%k)
 111         if k in self.d:
 112             raise InvalidStatusFile("Multiple occurences of key %r"%k)
 113         self.d[k]=v
 114
 115     def do_read(self, lines):
 116         """
 117         Internal: Parse a status file given as list/iterator of lines
 118         """
 119
 120         stk = [] # stack
 121         this = self # currently read nested Status
 122         lastk = None # for multiline appending
 123         for x in lines:
 124             x = x.rstrip("\n").lstrip(" \t")
 125             if x=="" or x.startswith("#"):
 126                 lastk = None
 127             else:
 128                 sep = x.find(":")
 129                 if sep > 0: # key:value
 130                     k = x[:sep].rstrip(" \t")
 131                     v = x[sep+1:]
 132                     this.read_val(k, v)
 133                     lastk = k
 134                 elif sep == 0: # continuation of multiline :value
 135                     if not lastk:
 136                         raise InvalidStatusFile("Parse error: key expected before ':'")
 137                     v = x[sep+1:]
 138                     this[lastk] += '\n' + v
 139                 elif x.endswith("("): # close subtree
 140                     k = x[:-1].rstrip(" \t")
 141                     new = Status()
 142                     this.read_val(k, new)
 143                     stk.append(this)
 144                     this = new
 145                     lastk = None
 146                 elif x == ")":
 147                     lastk = None
 148                     if len(stk) == 0:
 149                         raise InvalidStatusFile("Parse error: incorrect nesting")
 150                     else:
 151                         this = stk.pop()
 152                 else:
 153                     raise InvalidStatusFile("Parse error: malformed line")
 154         if stk:
 155             raise InvalidStatusFile("Parse error: not all subtrees closed")