ceph/src/ceph-volume/ceph_volume/decorators.py

   1 import os
   2 import sys
   3 from ceph_volume import terminal, exceptions
   4 from functools import wraps
   5
   6
   7 def needs_root(func):
   8     """
   9     Check for super user privileges on functions/methods. Raise
  10     ``SuperUserError`` with a nice message.
  11     """
  12     @wraps(func)
  13     def is_root(*a, **kw):
  14         if not os.getuid() == 0:
  15             raise exceptions.SuperUserError()
  16         return func(*a, **kw)
  17     return is_root
  18
  19
  20 def catches(catch=None, handler=None, exit=True):
  21     """
  22     Very simple decorator that tries any of the exception(s) passed in as
  23     a single exception class or tuple (containing multiple ones) returning the
  24     exception message and optionally handling the problem if it rises with the
  25     handler if it is provided.
  26
  27     So instead of douing something like this::
  28
  29         def bar():
  30             try:
  31                 some_call()
  32                 print "Success!"
  33             except TypeError, exc:
  34                 print "Error while handling some call: %s" % exc
  35                 sys.exit(1)
  36
  37     You would need to decorate it like this to have the same effect::
  38
  39         @catches(TypeError)
  40         def bar():
  41             some_call()
  42             print "Success!"
  43
  44     If multiple exceptions need to be catched they need to be provided as a
  45     tuple::
  46
  47         @catches((TypeError, AttributeError))
  48         def bar():
  49             some_call()
  50             print "Success!"
  51     """
  52     catch = catch or Exception
  53
  54     def decorate(f):
  55
  56         @wraps(f)
  57         def newfunc(*a, **kw):
  58             try:
  59                 return f(*a, **kw)
  60             except catch as e:
  61                 if os.environ.get('CEPH_VOLUME_DEBUG'):
  62                     raise
  63                 if handler:
  64                     return handler(e)
  65                 else:
  66                     sys.stderr.write(make_exception_message(e))
  67                     if exit:
  68                         sys.exit(1)
  69         return newfunc
  70
  71     return decorate
  72
  73 #
  74 # Decorator helpers
  75 #
  76
  77
  78 def make_exception_message(exc):
  79     """
  80     An exception is passed in and this function
  81     returns the proper string depending on the result
  82     so it is readable enough.
  83     """
  84     if str(exc):
  85         return '%s %s: %s\n' % (terminal.red_arrow, exc.__class__.__name__, exc)
  86     else:
  87         return '%s %s\n' % (terminal.red_arrow, exc.__class__.__name__)