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     sys.stdout.flush()
  23     _hard_write(sys.stderr.fileno(), s)
  24
  25
  26 def mkdirp(d):
  27     try:
  28         os.makedirs(d)
  29     except OSError, e:
  30         if e.errno == errno.EEXIST:
  31             pass
  32         else:
  33             raise
  34
  35
  36 def next(it):
  37     try:
  38         return it.next()
  39     except StopIteration:
  40         return None
  41
  42
  43 def unlink(f):
  44     try:
  45         os.unlink(f)
  46     except OSError, e:
  47         if e.errno == errno.ENOENT:
  48             pass  # it doesn't exist, that's what you asked for
  49
  50
  51 def readpipe(argv):
  52     p = subprocess.Popen(argv, stdout=subprocess.PIPE)
  53     r = p.stdout.read()
  54     p.wait()
  55     return r
  56
  57
  58 # FIXME: this function isn't very generic, because it splits the filename
  59 # in an odd way and depends on a terminating '/' to indicate directories.
  60 # But it's used in a couple of places, so let's put it here.
  61 def pathsplit(p):
  62     l = p.split('/')
  63     l = [i+'/' for i in l[:-1]] + l[-1:]
  64     if l[-1] == '':
  65         l.pop()  # extra blank caused by terminating '/'
  66     return l
  67
  68
  69 # like os.path.realpath, but doesn't follow a symlink for the last element.
  70 # (ie. if 'p' itself is itself a symlink, this one won't follow it)
  71 def realpath(p):
  72     try:
  73         st = os.lstat(p)
  74     except OSError:
  75         st = None
  76     if st and stat.S_ISLNK(st.st_mode):
  77         (dir, name) = os.path.split(p)
  78         dir = os.path.realpath(dir)
  79         out = os.path.join(dir, name)
  80     else:
  81         out = os.path.realpath(p)
  82     #log('realpathing:%r,%r\n' % (p, out))
  83     return out
  84
  85
  86 _username = None
  87 def username():
  88     global _username
  89     if not _username:
  90         uid = os.getuid()
  91         try:
  92             _username = pwd.getpwuid(uid)[0]
  93         except KeyError:
  94             _username = 'user%d' % uid
  95     return _username
  96
  97
  98 _userfullname = None
  99 def userfullname():
 100     global _userfullname
 101     if not _userfullname:
 102         uid = os.getuid()
 103         try:
 104             _userfullname = pwd.getpwuid(uid)[4].split(',')[0]
 105         except KeyError:
 106             _userfullname = 'user%d' % uid
 107     return _userfullname
 108
 109
 110 _hostname = None
 111 def hostname():
 112     global _hostname
 113     if not _hostname:
 114         _hostname = socket.getfqdn()
 115     return _hostname
 116
 117
 118 class NotOk(Exception):
 119     pass
 120
 121 class Conn:
 122     def __init__(self, inp, outp):
 123         self.inp = inp
 124         self.outp = outp
 125
 126     def read(self, size):
 127         self.outp.flush()
 128         return self.inp.read(size)
 129
 130     def readline(self):
 131         self.outp.flush()
 132         return self.inp.readline()
 133
 134     def write(self, data):
 135         #log('%d writing: %d bytes\n' % (os.getpid(), len(data)))
 136         self.outp.write(data)
 137
 138     def has_input(self):
 139         [rl, wl, xl] = select.select([self.inp.fileno()], [], [], 0)
 140         if rl:
 141             assert(rl[0] == self.inp.fileno())
 142             return True
 143         else:
 144             return None
 145
 146     def ok(self):
 147         self.write('\nok\n')
 148
 149     def error(self, s):
 150         s = re.sub(r'\s+', ' ', str(s))
 151         self.write('\nerror %s\n' % s)
 152
 153     def _check_ok(self, onempty):
 154         self.outp.flush()
 155         rl = ''
 156         for rl in linereader(self.inp):
 157             #log('%d got line: %r\n' % (os.getpid(), rl))
 158             if not rl:  # empty line
 159                 continue
 160             elif rl == 'ok':
 161                 return None
 162             elif rl.startswith('error '):
 163                 #log('client: error: %s\n' % rl[6:])
 164                 return NotOk(rl[6:])
 165             else:
 166                 onempty(rl)
 167         raise Exception('server exited unexpectedly; see errors above')
 168
 169     def drain_and_check_ok(self):
 170         def onempty(rl):
 171             pass
 172         return self._check_ok(onempty)
 173
 174     def check_ok(self):
 175         def onempty(rl):
 176             raise Exception('expected "ok", got %r' % rl)
 177         return self._check_ok(onempty)
 178
 179
 180 def linereader(f):
 181     while 1:
 182         line = f.readline()
 183         if not line:
 184             break
 185         yield line[:-1]
 186
 187
 188 def chunkyreader(f, count = None):
 189     if count != None:
 190         while count > 0:
 191             b = f.read(min(count, 65536))
 192             if not b:
 193                 raise IOError('EOF with %d bytes remaining' % count)
 194             yield b
 195             count -= len(b)
 196     else:
 197         while 1:
 198             b = f.read(65536)
 199             if not b: break
 200             yield b
 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 def atoi(s):
 255     try:
 256         return int(s or '0')
 257     except ValueError:
 258         return 0
 259
 260
 261 saved_errors = []
 262 def add_error(e):
 263     saved_errors.append(e)
 264     log('%-70s\n' % e)
 265
 266 istty = os.isatty(2) or atoi(os.environ.get('BUP_FORCE_TTY'))
 267 def progress(s):
 268     if istty:
 269         log(s)
 270
 271
 272 def handle_ctrl_c():
 273     oldhook = sys.excepthook
 274     def newhook(exctype, value, traceback):
 275         if exctype == KeyboardInterrupt:
 276             log('Interrupted.\n')
 277         else:
 278             return oldhook(exctype, value, traceback)
 279     sys.excepthook = newhook
 280
 281
 282 def columnate(l, prefix):
 283     l = l[:]
 284     clen = max(len(s) for s in l)
 285     ncols = (78 - len(prefix)) / (clen + 2)
 286     if ncols <= 1:
 287         ncols = 1
 288         clen = 0
 289     cols = []
 290     while len(l) % ncols:
 291         l.append('')
 292     rows = len(l)/ncols
 293     for s in range(0, len(l), rows):
 294         cols.append(l[s:s+rows])
 295     out = ''
 296     for row in zip(*cols):
 297         out += prefix + ''.join(('%-*s' % (clen+2, s)) for s in row) + '\n'
 298     return out
 299
 300
 301 # hashlib is only available in python 2.5 or higher, but the 'sha' module
 302 # produces a DeprecationWarning in python 2.6 or higher.  We want to support
 303 # python 2.4 and above without any stupid warnings, so let's try using hashlib
 304 # first, and downgrade if it fails.
 305 try:
 306     import hashlib
 307 except ImportError:
 308     import sha
 309     Sha1 = sha.sha
 310 else:
 311     Sha1 = hashlib.sha1