vstf/vstf/common/decorator.py

   1 ##############################################################################
   2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
   3 #
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9
  10 _DEFAULTS = "vstf check key defaults".encode()
  11
  12
  13 def check(key, choices=[], defaults=_DEFAULTS):
  14     def _deco(func):
  15         def __deco(*args, **kwargs):
  16             if key not in kwargs:
  17                 if defaults != _DEFAULTS:
  18                     kwargs[key] = defaults
  19                 else:
  20                     raise Exception("Error: '%s' is needed in %s" % (key, func))
  21
  22             if choices and kwargs[key] not in choices:
  23                 raise Exception("Error: %s :%s" % (key, kwargs[key]))
  24             ret = func(*args, **kwargs)
  25             return ret
  26
  27         return __deco
  28
  29     return _deco
  30
  31
  32 def dcheck(key, choices=[]):
  33     def _deco(func):
  34         def __deco(*args):
  35             if len(args) == 2:
  36                 values = args[1]
  37             elif len(args) == 1:
  38                 values = args[0]
  39             else:
  40                 values = None
  41             if isinstance(values, dict):
  42                 if key not in values:
  43                     raise Exception("Error: '%s' is needed in %s" % (key, func))
  44                 if choices and values[key] not in choices:
  45                     raise Exception("Error: %s :%s" % (key, values[key]))
  46             ret = func(*args)
  47             return ret
  48
  49         return __deco
  50
  51     return _deco
  52
  53
  54 def vstf_input(key, types=str, choices=[], default=None):
  55     def _deco(func):
  56         def __deco(*args):
  57             ret = func(*args)
  58             if not isinstance(ret, dict):
  59                 ret = {}
  60             in_str = "----> %s : " % key
  61             if choices:
  62                 in_str = "---- %s\n" % (str(choices)) + in_str
  63             while True:
  64                 if types == list or types == dict:
  65                     value = input(in_str)
  66                 else:
  67                     value = raw_input(in_str)
  68                     value = types(value)
  69                 if not value:
  70                     value = default
  71                 if not choices or value in choices:
  72                     break
  73             ret.update({key: value})
  74             return ret
  75
  76         return __deco
  77
  78     return _deco
  79
  80
  81 def namespace():
  82     def _deco(func):
  83         def __deco(*args, **kwargs):
  84             ret = func(*args, **kwargs)
  85             nspace = kwargs.get("namespace", None)
  86             if nspace:
  87                 ret = "ip netns exec %(namespace)s " % {"namespace": nspace} + ret
  88             return ret
  89
  90         return __deco
  91
  92     return _deco