vstf/vstf/common/decorator.py

   1 #!/usr/bin/python
   2 # -*- coding: utf8 -*-
   3 # author: wly
   4 # date: 2015-09-09
   5 # see license for license details
   6 _DEFAULTS = "vstf check key defaults".encode()
   7
   8
   9 def check(key, choices=[], defaults=_DEFAULTS):
  10     def _deco(func):
  11         def __deco(*args, **kwargs):
  12             if key not in kwargs:
  13                 if defaults != _DEFAULTS:
  14                     kwargs[key] = defaults
  15                 else:
  16                     raise Exception("Error: '%s' is needed in %s" % (key, func))
  17
  18             if choices and kwargs[key] not in choices:
  19                 raise Exception("Error: %s :%s" % (key, kwargs[key]))
  20             ret = func(*args, **kwargs)
  21             return ret
  22
  23         return __deco
  24
  25     return _deco
  26
  27
  28 def dcheck(key, choices=[]):
  29     def _deco(func):
  30         def __deco(*args):
  31             if len(args) == 2:
  32                 values = args[1]
  33             elif len(args) == 1:
  34                 values = args[0]
  35             else:
  36                 values = None
  37             if isinstance(values, dict):
  38                 if key not in values:
  39                     raise Exception("Error: '%s' is needed in %s" % (key, func))
  40                 if choices and values[key] not in choices:
  41                     raise Exception("Error: %s :%s" % (key, values[key]))
  42             ret = func(*args)
  43             return ret
  44
  45         return __deco
  46
  47     return _deco
  48
  49
  50 def vstf_input(key, types=str, choices=[], default=None):
  51     def _deco(func):
  52         def __deco(*args):
  53             ret = func(*args)
  54             if not isinstance(ret, dict):
  55                 ret = {}
  56             in_str = "----> %s : " % key
  57             if choices:
  58                 in_str = "---- %s\n" % (str(choices)) + in_str
  59             while True:
  60                 if types == list or types == dict:
  61                     value = input(in_str)
  62                 else:
  63                     value = raw_input(in_str)
  64                     value = types(value)
  65                 if not value:
  66                     value = default
  67                 if not choices or value in choices:
  68                     break
  69             ret.update({key: value})
  70             return ret
  71
  72         return __deco
  73
  74     return _deco
  75
  76
  77 def namespace():
  78     def _deco(func):
  79         def __deco(*args, **kwargs):
  80             ret = func(*args, **kwargs)
  81             nspace = kwargs.get("namespace", None)
  82             if nspace:
  83                 ret = "ip netns exec %(namespace)s " % {"namespace": nspace} + ret
  84             return ret
  85
  86         return __deco
  87
  88     return _deco