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