lib/bup/helpers.py

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