helpers.py

   1 import sys, os, pwd, subprocess, errno, socket, select, mmap, stat, re
   2
   3
   4 def log(s):
   5     sys.stderr.write(s)
   6
   7
   8 def mkdirp(d):
   9     try:
  10         os.makedirs(d)
  11     except OSError, e:
  12         if e.errno == errno.EEXIST:
  13             pass
  14         else:
  15             raise
  16
  17
  18 def next(it):
  19     try:
  20         return it.next()
  21     except StopIteration:
  22         return None
  23
  24
  25 def unlink(f):
  26     try:
  27         os.unlink(f)
  28     except OSError, e:
  29         if e.errno == errno.ENOENT:
  30             pass  # it doesn't exist, that's what you asked for
  31
  32
  33 def readpipe(argv):
  34     p = subprocess.Popen(argv, stdout=subprocess.PIPE)
  35     r = p.stdout.read()
  36     p.wait()
  37     return r
  38
  39
  40 # FIXME: this function isn't very generic, because it splits the filename
  41 # in an odd way and depends on a terminating '/' to indicate directories.
  42 # But it's used in a couple of places, so let's put it here.
  43 def pathsplit(p):
  44     l = p.split('/')
  45     l = [i+'/' for i in l[:-1]] + l[-1:]
  46     if l[-1] == '':
  47         l.pop()  # extra blank caused by terminating '/'
  48     return l
  49
  50
  51 # like os.path.realpath, but doesn't follow a symlink for the last element.
  52 # (ie. if 'p' itself is itself a symlink, this one won't follow it)
  53 def realpath(p):
  54     try:
  55         st = os.lstat(p)
  56     except OSError:
  57         st = None
  58     if st and stat.S_ISLNK(st.st_mode):
  59         (dir, name) = os.path.split(p)
  60         dir = os.path.realpath(dir)
  61         out = os.path.join(dir, name)
  62     else:
  63         out = os.path.realpath(p)
  64     #log('realpathing:%r,%r\n' % (p, out))
  65     return out
  66
  67
  68 _username = None
  69 def username():
  70     global _username
  71     if not _username:
  72         uid = os.getuid()
  73         try:
  74             _username = pwd.getpwuid(uid)[0]
  75         except KeyError:
  76             _username = 'user%d' % uid
  77     return _username
  78
  79
  80 _userfullname = None
  81 def userfullname():
  82     global _userfullname
  83     if not _userfullname:
  84         uid = os.getuid()
  85         try:
  86             _userfullname = pwd.getpwuid(uid)[4].split(',')[0]
  87         except KeyError:
  88             _userfullname = 'user%d' % uid
  89     return _userfullname
  90
  91
  92 _hostname = None
  93 def hostname():
  94     global _hostname
  95     if not _hostname:
  96         _hostname = socket.getfqdn()
  97     return _hostname
  98
  99
 100 class NotOk(Exception):
 101     pass
 102
 103 class Conn:
 104     def __init__(self, inp, outp):
 105         self.inp = inp
 106         self.outp = outp
 107
 108     def read(self, size):
 109         self.outp.flush()
 110         return self.inp.read(size)
 111
 112     def readline(self):
 113         self.outp.flush()
 114         return self.inp.readline()
 115
 116     def write(self, data):
 117         #log('%d writing: %d bytes\n' % (os.getpid(), len(data)))
 118         self.outp.write(data)
 119
 120     def has_input(self):
 121         [rl, wl, xl] = select.select([self.inp.fileno()], [], [], 0)
 122         if rl:
 123             assert(rl[0] == self.inp.fileno())
 124             return True
 125         else:
 126             return None
 127
 128     def ok(self):
 129         self.write('\nok\n')
 130
 131     def error(self, s):
 132         s = re.sub(r'\s+', ' ', str(s))
 133         self.write('\nerror %s\n' % s)
 134
 135     def _check_ok(self, onempty):
 136         self.outp.flush()
 137         rl = ''
 138         for rl in linereader(self.inp):
 139             #log('%d got line: %r\n' % (os.getpid(), rl))
 140             if not rl:  # empty line
 141                 continue
 142             elif rl == 'ok':
 143                 return None
 144             elif rl.startswith('error '):
 145                 #log('client: error: %s\n' % rl[6:])
 146                 return NotOk(rl[6:])
 147             else:
 148                 onempty(rl)
 149         raise Exception('server exited unexpectedly; see errors above')
 150
 151     def drain_and_check_ok(self):
 152         def onempty(rl):
 153             pass
 154         return self._check_ok(onempty)
 155
 156     def check_ok(self):
 157         def onempty(rl):
 158             raise Exception('expected "ok", got %r' % rl)
 159         return self._check_ok(onempty)
 160
 161
 162 def linereader(f):
 163     while 1:
 164         line = f.readline()
 165         if not line:
 166             break
 167         yield line[:-1]
 168
 169
 170 def chunkyreader(f, count = None):
 171     if count != None:
 172         while count > 0:
 173             b = f.read(min(count, 65536))
 174             if not b:
 175                 raise IOError('EOF with %d bytes remaining' % count)
 176             yield b
 177             count -= len(b)
 178     else:
 179         while 1:
 180             b = f.read(65536)
 181             if not b: break
 182             yield b
 183
 184
 185 class AutoFlushIter:
 186     def __init__(self, it, ondone = None):
 187         self.it = it
 188         self.ondone = ondone
 189
 190     def __iter__(self):
 191         return self
 192
 193     def next(self):
 194         return self.it.next()
 195
 196     def __del__(self):
 197         for i in self.it:
 198             pass
 199         if self.ondone:
 200             self.ondone()
 201
 202
 203 def slashappend(s):
 204     if s and not s.endswith('/'):
 205         return s + '/'
 206     else:
 207         return s
 208
 209
 210 def _mmap_do(f, len, flags, prot):
 211     if not len:
 212         st = os.fstat(f.fileno())
 213         len = st.st_size
 214     map = mmap.mmap(f.fileno(), len, flags, prot)
 215     f.close()  # map will persist beyond file close
 216     return map
 217
 218
 219 def mmap_read(f, len = 0):
 220     return _mmap_do(f, len, mmap.MAP_PRIVATE, mmap.PROT_READ)
 221
 222
 223 def mmap_readwrite(f, len = 0):
 224     return _mmap_do(f, len, mmap.MAP_SHARED, mmap.PROT_READ|mmap.PROT_WRITE)
 225
 226
 227 def parse_num(s):
 228     g = re.match(r'([-+\d.e]+)\s*(\w*)', str(s))
 229     if not g:
 230         raise ValueError("can't parse %r as a number" % s)
 231     (val, unit) = g.groups()
 232     num = float(val)
 233     unit = unit.lower()
 234     if unit in ['t', 'tb']:
 235         mult = 1024*1024*1024*1024
 236     elif unit in ['g', 'gb']:
 237         mult = 1024*1024*1024
 238     elif unit in ['m', 'mb']:
 239         mult = 1024*1024
 240     elif unit in ['k', 'kb']:
 241         mult = 1024
 242     elif unit in ['', 'b']:
 243         mult = 1
 244     else:
 245         raise ValueError("invalid unit %r in number %r" % (unit, s))
 246     return int(num*mult)
 247
 248
 249 # count the number of elements in an iterator (consumes the iterator)
 250 def count(l):
 251     return reduce(lambda x,y: x+1, l)
 252
 253
 254 saved_errors = []
 255 def add_error(e):
 256     saved_errors.append(e)
 257     log('%-70s\n' % e)
 258
 259
 260 istty = os.isatty(2)
 261 def progress(s):
 262     if istty:
 263         log(s)