helpers.py

   1 import sys, os, pwd, subprocess, errno, socket, select, mmap, stat
   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 Conn:
 101     def __init__(self, inp, outp):
 102         self.inp = inp
 103         self.outp = outp
 104
 105     def read(self, size):
 106         self.outp.flush()
 107         return self.inp.read(size)
 108
 109     def readline(self):
 110         self.outp.flush()
 111         return self.inp.readline()
 112
 113     def write(self, data):
 114         #log('%d writing: %d bytes\n' % (os.getpid(), len(data)))
 115         self.outp.write(data)
 116
 117     def has_input(self):
 118         [rl, wl, xl] = select.select([self.inp.fileno()], [], [], 0)
 119         if rl:
 120             assert(rl[0] == self.inp.fileno())
 121             return True
 122         else:
 123             return None
 124
 125     def ok(self):
 126         self.write('\nok\n')
 127
 128     def drain_and_check_ok(self):
 129         self.outp.flush()
 130         rl = ''
 131         for rl in linereader(self.inp):
 132             #log('%d got line: %r\n' % (os.getpid(), rl))
 133             if not rl:  # empty line
 134                 continue
 135             elif rl == 'ok':
 136                 return True
 137             else:
 138                 pass # ignore line
 139         # NOTREACHED
 140
 141     def check_ok(self):
 142         self.outp.flush()
 143         rl = ''
 144         for rl in linereader(self.inp):
 145             #log('%d got line: %r\n' % (os.getpid(), rl))
 146             if not rl:  # empty line
 147                 continue
 148             elif rl == 'ok':
 149                 return True
 150             else:
 151                 raise Exception('expected "ok", got %r' % rl)
 152         raise Exception('server exited unexpectedly; see errors above')
 153
 154
 155 def linereader(f):
 156     while 1:
 157         line = f.readline()
 158         if not line:
 159             break
 160         yield line[:-1]
 161
 162
 163 def chunkyreader(f, count = None):
 164     if count != None:
 165         while count > 0:
 166             b = f.read(min(count, 65536))
 167             if not b:
 168                 raise IOError('EOF with %d bytes remaining' % count)
 169             yield b
 170             count -= len(b)
 171     else:
 172         while 1:
 173             b = f.read(65536)
 174             if not b: break
 175             yield b
 176
 177
 178 def slashappend(s):
 179     if s and not s.endswith('/'):
 180         return s + '/'
 181     else:
 182         return s
 183
 184
 185 def _mmap_do(f, len, flags, prot):
 186     if not len:
 187         st = os.fstat(f.fileno())
 188         len = st.st_size
 189     map = mmap.mmap(f.fileno(), len, flags, prot)
 190     f.close()  # map will persist beyond file close
 191     return map
 192
 193
 194 def mmap_read(f, len = 0):
 195     return _mmap_do(f, len, mmap.MAP_PRIVATE, mmap.PROT_READ)
 196
 197
 198 def mmap_readwrite(f, len = 0):
 199     return _mmap_do(f, len, mmap.MAP_SHARED, mmap.PROT_READ|mmap.PROT_WRITE)
 200
 201
 202 # count the number of elements in an iterator (consumes the iterator)
 203 def count(l):
 204     return reduce(lambda x,y: x+1, l)
 205
 206
 207 saved_errors = []
 208 def add_error(e):
 209     saved_errors.append(e)
 210     log('%-70s\n' % e)
 211
 212
 213 istty = os.isatty(2)
 214 def progress(s):
 215     if istty:
 216         log(s)