t/moe/status.py

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