lib/bup/helpers.py

   1 """Helper functions and classes for bup."""
   2 import sys, os, pwd, subprocess, errno, socket, select, mmap, stat, re
   3 from bup import _version
   4
   5
   6 # Write (blockingly) to sockets that may or may not be in blocking mode.
   7 # We need this because our stderr is sometimes eaten by subprocesses
   8 # (probably ssh) that sometimes make it nonblocking, if only temporarily,
   9 # leading to race conditions.  Ick.  We'll do it the hard way.
  10 def _hard_write(fd, buf):
  11     while buf:
  12         (r,w,x) = select.select([], [fd], [], None)
  13         if not w:
  14             raise IOError('select(fd) returned without being writable')
  15         try:
  16             sz = os.write(fd, buf)
  17         except OSError, e:
  18             if e.errno != errno.EAGAIN:
  19                 raise
  20         assert(sz >= 0)
  21         buf = buf[sz:]
  22
  23 def log(s):
  24     """Print a log message to stderr."""
  25     sys.stdout.flush()
  26     _hard_write(sys.stderr.fileno(), s)
  27
  28
  29 def mkdirp(d, mode=None):
  30     """Recursively create directories on path 'd'.
  31
  32     Unlike os.makedirs(), it doesn't raise an exception if the last element of
  33     the path already exists.
  34     """
  35     try:
  36         if mode:
  37             os.makedirs(d, mode)
  38         else:
  39             os.makedirs(d)
  40     except OSError, e:
  41         if e.errno == errno.EEXIST:
  42             pass
  43         else:
  44             raise
  45
  46
  47 def next(it):
  48     """Get the next item from an iterator, None if we reached the end."""
  49     try:
  50         return it.next()
  51     except StopIteration:
  52         return None
  53
  54
  55 def unlink(f):
  56     """Delete a file at path 'f' if it currently exists.
  57
  58     Unlike os.unlink(), does not throw an exception if the file didn't already
  59     exist.
  60     """
  61     try:
  62         os.unlink(f)
  63     except OSError, e:
  64         if e.errno == errno.ENOENT:
  65             pass  # it doesn't exist, that's what you asked for
  66
  67
  68 def readpipe(argv):
  69     """Run a subprocess and return its output."""
  70     p = subprocess.Popen(argv, stdout=subprocess.PIPE)
  71     r = p.stdout.read()
  72     p.wait()
  73     return r
  74
  75
  76 def realpath(p):
  77     """Get the absolute path of a file.
  78
  79     Behaves like os.path.realpath, but doesn't follow a symlink for the last
  80     element. (ie. if 'p' itself is a symlink, this one won't follow it, but it
  81     will follow symlinks in p's directory)
  82     """
  83     try:
  84         st = os.lstat(p)
  85     except OSError:
  86         st = None
  87     if st and stat.S_ISLNK(st.st_mode):
  88         (dir, name) = os.path.split(p)
  89         dir = os.path.realpath(dir)
  90         out = os.path.join(dir, name)
  91     else:
  92         out = os.path.realpath(p)
  93     #log('realpathing:%r,%r\n' % (p, out))
  94     return out
  95
  96
  97 _username = None
  98 def username():
  99     """Get the user's login name."""
 100     global _username
 101     if not _username:
 102         uid = os.getuid()
 103         try:
 104             _username = pwd.getpwuid(uid)[0]
 105         except KeyError:
 106             _username = 'user%d' % uid
 107     return _username
 108
 109
 110 _userfullname = None
 111 def userfullname():
 112     """Get the user's full name."""
 113     global _userfullname
 114     if not _userfullname:
 115         uid = os.getuid()
 116         try:
 117             _userfullname = pwd.getpwuid(uid)[4].split(',')[0]
 118         except KeyError:
 119             _userfullname = 'user%d' % uid
 120     return _userfullname
 121
 122
 123 _hostname = None
 124 def hostname():
 125     """Get the FQDN of this machine."""
 126     global _hostname
 127     if not _hostname:
 128         _hostname = socket.getfqdn()
 129     return _hostname
 130
 131
 132 _resource_path = None
 133 def resource_path(subdir=''):
 134     global _resource_path
 135     if not _resource_path:
 136         _resource_path = os.environ.get('BUP_RESOURCE_PATH') or '.'
 137     return os.path.join(_resource_path, subdir)
 138
 139 class NotOk(Exception):
 140     pass
 141
 142 class Conn:
 143     """A helper class for bup's client-server protocol."""
 144     def __init__(self, inp, outp):
 145         self.inp = inp
 146         self.outp = outp
 147
 148     def read(self, size):
 149         """Read 'size' bytes from input stream."""
 150         self.outp.flush()
 151         return self.inp.read(size)
 152
 153     def readline(self):
 154         """Read from input stream until a newline is found."""
 155         self.outp.flush()
 156         return self.inp.readline()
 157
 158     def write(self, data):
 159         """Write 'data' to output stream."""
 160         #log('%d writing: %d bytes\n' % (os.getpid(), len(data)))
 161         self.outp.write(data)
 162
 163     def has_input(self):
 164         """Return true if input stream is readable."""
 165         [rl, wl, xl] = select.select([self.inp.fileno()], [], [], 0)
 166         if rl:
 167             assert(rl[0] == self.inp.fileno())
 168             return True
 169         else:
 170             return None
 171
 172     def ok(self):
 173         """Indicate end of output from last sent command."""
 174         self.write('\nok\n')
 175
 176     def error(self, s):
 177         """Indicate server error to the client."""
 178         s = re.sub(r'\s+', ' ', str(s))
 179         self.write('\nerror %s\n' % s)
 180
 181     def _check_ok(self, onempty):
 182         self.outp.flush()
 183         rl = ''
 184         for rl in linereader(self.inp):
 185             #log('%d got line: %r\n' % (os.getpid(), rl))
 186             if not rl:  # empty line
 187                 continue
 188             elif rl == 'ok':
 189                 return None
 190             elif rl.startswith('error '):
 191                 #log('client: error: %s\n' % rl[6:])
 192                 return NotOk(rl[6:])
 193             else:
 194                 onempty(rl)
 195         raise Exception('server exited unexpectedly; see errors above')
 196
 197     def drain_and_check_ok(self):
 198         """Remove all data for the current command from input stream."""
 199         def onempty(rl):
 200             pass
 201         return self._check_ok(onempty)
 202
 203     def check_ok(self):
 204         """Verify that server action completed successfully."""
 205         def onempty(rl):
 206             raise Exception('expected "ok", got %r' % rl)
 207         return self._check_ok(onempty)
 208
 209
 210 def linereader(f):
 211     """Generate a list of input lines from 'f' without terminating newlines."""
 212     while 1:
 213         line = f.readline()
 214         if not line:
 215             break
 216         yield line[:-1]
 217
 218
 219 def chunkyreader(f, count = None):
 220     """Generate a list of chunks of data read from 'f'.
 221
 222     If count is None, read until EOF is reached.
 223
 224     If count is a positive integer, read 'count' bytes from 'f'. If EOF is
 225     reached while reading, raise IOError.
 226     """
 227     if count != None:
 228         while count > 0:
 229             b = f.read(min(count, 65536))
 230             if not b:
 231                 raise IOError('EOF with %d bytes remaining' % count)
 232             yield b
 233             count -= len(b)
 234     else:
 235         while 1:
 236             b = f.read(65536)
 237             if not b: break
 238             yield b
 239
 240
 241 def slashappend(s):
 242     """Append "/" to 's' if it doesn't aleady end in "/"."""
 243     if s and not s.endswith('/'):
 244         return s + '/'
 245     else:
 246         return s
 247
 248
 249 def _mmap_do(f, sz, flags, prot):
 250     if not sz:
 251         st = os.fstat(f.fileno())
 252         sz = st.st_size
 253     map = mmap.mmap(f.fileno(), sz, flags, prot)
 254     f.close()  # map will persist beyond file close
 255     return map
 256
 257
 258 def mmap_read(f, sz = 0):
 259     """Create a read-only memory mapped region on file 'f'.
 260
 261     If sz is 0, the region will cover the entire file.
 262     """
 263     return _mmap_do(f, sz, mmap.MAP_PRIVATE, mmap.PROT_READ)
 264
 265
 266 def mmap_readwrite(f, sz = 0):
 267     """Create a read-write memory mapped region on file 'f'.
 268
 269     If sz is 0, the region will cover the entire file.
 270     """
 271     return _mmap_do(f, sz, mmap.MAP_SHARED, mmap.PROT_READ|mmap.PROT_WRITE)
 272
 273
 274 def parse_num(s):
 275     """Parse data size information into a float number.
 276
 277     Here are some examples of conversions:
 278         199.2k means 203981 bytes
 279         1GB means 1073741824 bytes
 280         2.1 tb means 2199023255552 bytes
 281     """
 282     g = re.match(r'([-+\d.e]+)\s*(\w*)', str(s))
 283     if not g:
 284         raise ValueError("can't parse %r as a number" % s)
 285     (val, unit) = g.groups()
 286     num = float(val)
 287     unit = unit.lower()
 288     if unit in ['t', 'tb']:
 289         mult = 1024*1024*1024*1024
 290     elif unit in ['g', 'gb']:
 291         mult = 1024*1024*1024
 292     elif unit in ['m', 'mb']:
 293         mult = 1024*1024
 294     elif unit in ['k', 'kb']:
 295         mult = 1024
 296     elif unit in ['', 'b']:
 297         mult = 1
 298     else:
 299         raise ValueError("invalid unit %r in number %r" % (unit, s))
 300     return int(num*mult)
 301
 302
 303 def count(l):
 304     """Count the number of elements in an iterator. (consumes the iterator)"""
 305     return reduce(lambda x,y: x+1, l)
 306
 307
 308 def atoi(s):
 309     """Convert the string 's' to an integer. Return 0 if s is not a number."""
 310     try:
 311         return int(s or '0')
 312     except ValueError:
 313         return 0
 314
 315
 316 saved_errors = []
 317 def add_error(e):
 318     """Append an error message to the list of saved errors.
 319
 320     Once processing is able to stop and output the errors, the saved errors are
 321     accessible in the module variable helpers.saved_errors.
 322     """
 323     saved_errors.append(e)
 324     log('%-70s\n' % e)
 325
 326 istty = os.isatty(2) or atoi(os.environ.get('BUP_FORCE_TTY'))
 327 def progress(s):
 328     """Calls log(s) if stderr is a TTY.  Does nothing otherwise."""
 329     if istty:
 330         log(s)
 331
 332
 333 def handle_ctrl_c():
 334     """Replace the default exception handler for KeyboardInterrupt (Ctrl-C).
 335
 336     The new exception handler will make sure that bup will exit without an ugly
 337     stacktrace when Ctrl-C is hit.
 338     """
 339     oldhook = sys.excepthook
 340     def newhook(exctype, value, traceback):
 341         if exctype == KeyboardInterrupt:
 342             log('Interrupted.\n')
 343         else:
 344             return oldhook(exctype, value, traceback)
 345     sys.excepthook = newhook
 346
 347
 348 def columnate(l, prefix):
 349     """Format elements of 'l' in columns with 'prefix' leading each line.
 350
 351     The number of columns is determined automatically based on the string
 352     lengths.
 353     """
 354     if not l:
 355         return ""
 356     l = l[:]
 357     clen = max(len(s) for s in l)
 358     ncols = (78 - len(prefix)) / (clen + 2)
 359     if ncols <= 1:
 360         ncols = 1
 361         clen = 0
 362     cols = []
 363     while len(l) % ncols:
 364         l.append('')
 365     rows = len(l)/ncols
 366     for s in range(0, len(l), rows):
 367         cols.append(l[s:s+rows])
 368     out = ''
 369     for row in zip(*cols):
 370         out += prefix + ''.join(('%-*s' % (clen+2, s)) for s in row) + '\n'
 371     return out
 372
 373
 374 # hashlib is only available in python 2.5 or higher, but the 'sha' module
 375 # produces a DeprecationWarning in python 2.6 or higher.  We want to support
 376 # python 2.4 and above without any stupid warnings, so let's try using hashlib
 377 # first, and downgrade if it fails.
 378 try:
 379     import hashlib
 380 except ImportError:
 381     import sha
 382     Sha1 = sha.sha
 383 else:
 384     Sha1 = hashlib.sha1
 385
 386
 387 def version_date():
 388     """Format bup's version date string for output."""
 389     return _version.DATE.split(' ')[0]
 390
 391 def version_commit():
 392     """Get the commit hash of bup's current version."""
 393     return _version.COMMIT
 394
 395 def version_tag():
 396     """Format bup's version tag (the official version number).
 397
 398     When generated from a commit other than one pointed to with a tag, the
 399     returned string will be "unknown-" followed by the first seven positions of
 400     the commit hash.
 401     """
 402     names = _version.NAMES.strip()
 403     assert(names[0] == '(')
 404     assert(names[-1] == ')')
 405     names = names[1:-1]
 406     l = [n.strip() for n in names.split(',')]
 407     for n in l:
 408         if n.startswith('tag: bup-'):
 409             return n[9:]
 410     return 'unknown-%s' % _version.COMMIT[:7]