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