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