helpers.py

   1 import sys, os, pwd, subprocess, errno, socket, select, mmap
   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 readpipe(argv):
  19     p = subprocess.Popen(argv, stdout=subprocess.PIPE)
  20     r = p.stdout.read()
  21     p.wait()
  22     return r
  23
  24
  25 _username = None
  26 def username():
  27     global _username
  28     if not _username:
  29         uid = os.getuid()
  30         try:
  31             _username = pwd.getpwuid(uid)[0]
  32         except KeyError:
  33             _username = 'user%d' % uid
  34     return _username
  35
  36
  37 _userfullname = None
  38 def userfullname():
  39     global _userfullname
  40     if not _userfullname:
  41         uid = os.getuid()
  42         try:
  43             _userfullname = pwd.getpwuid(uid)[4].split(',')[0]
  44         except KeyError:
  45             _userfullname = 'user%d' % uid
  46     return _userfullname
  47
  48
  49 _hostname = None
  50 def hostname():
  51     global _hostname
  52     if not _hostname:
  53         _hostname = socket.getfqdn()
  54     return _hostname
  55
  56
  57 class Conn:
  58     def __init__(self, inp, outp):
  59         self.inp = inp
  60         self.outp = outp
  61
  62     def read(self, size):
  63         self.outp.flush()
  64         return self.inp.read(size)
  65
  66     def readline(self):
  67         self.outp.flush()
  68         return self.inp.readline()
  69
  70     def write(self, data):
  71         #log('%d writing: %d bytes\n' % (os.getpid(), len(data)))
  72         self.outp.write(data)
  73
  74     def has_input(self):
  75         [rl, wl, xl] = select.select([self.inp.fileno()], [], [], 0)
  76         if rl:
  77             assert(rl[0] == self.inp.fileno())
  78             return True
  79         else:
  80             return None
  81
  82     def ok(self):
  83         self.write('\nok\n')
  84
  85     def drain_and_check_ok(self):
  86         self.outp.flush()
  87         rl = ''
  88         for rl in linereader(self.inp):
  89             #log('%d got line: %r\n' % (os.getpid(), rl))
  90             if not rl:  # empty line
  91                 continue
  92             elif rl == 'ok':
  93                 return True
  94             else:
  95                 pass # ignore line
  96         # NOTREACHED
  97
  98     def check_ok(self):
  99         self.outp.flush()
 100         rl = ''
 101         for rl in linereader(self.inp):
 102             #log('%d got line: %r\n' % (os.getpid(), rl))
 103             if not rl:  # empty line
 104                 continue
 105             elif rl == 'ok':
 106                 return True
 107             else:
 108                 raise Exception('expected "ok", got %r' % rl)
 109         raise Exception('server exited unexpectedly; see errors above')
 110
 111
 112 def linereader(f):
 113     while 1:
 114         line = f.readline()
 115         if not line:
 116             break
 117         yield line[:-1]
 118
 119
 120 def chunkyreader(f, count = None):
 121     if count != None:
 122         while count > 0:
 123             b = f.read(min(count, 65536))
 124             if not b:
 125                 raise IOError('EOF with %d bytes remaining' % count)
 126             yield b
 127             count -= len(b)
 128     else:
 129         while 1:
 130             b = f.read(65536)
 131             if not b: break
 132             yield b
 133
 134
 135 def slashappend(s):
 136     if s and not s.endswith('/'):
 137         return s + '/'
 138     else:
 139         return s
 140
 141
 142 def _mmap_do(f, len, flags, prot):
 143     if not len:
 144         st = os.fstat(f.fileno())
 145         len = st.st_size
 146     map = mmap.mmap(f.fileno(), len, flags, prot)
 147     f.close()  # map will persist beyond file close
 148     return map
 149
 150
 151 def mmap_read(f, len = 0):
 152     return _mmap_do(f, len, mmap.MAP_PRIVATE, mmap.PROT_READ)
 153
 154
 155 def mmap_readwrite(f, len = 0):
 156     return _mmap_do(f, len, mmap.MAP_SHARED, mmap.PROT_READ|mmap.PROT_WRITE)