summaryrefslogtreecommitdiffstats
path: root/vstf/vstf/common/decorator.py
blob: cafbb3ffd2250fa9f365796912eb669bfc1ee946 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#!/usr/bin/python
# -*- coding: utf8 -*-
# author: wly
# date: 2015-09-09
# see license for license details
_DEFAULTS = "vstf check key defaults".encode()


def check(key, choices=[], defaults=_DEFAULTS):
    def _deco(func):
        def __deco(*args, **kwargs):
            if key not in kwargs:
                if defaults != _DEFAULTS:
                    kwargs[key] = defaults
                else:
                    raise Exception("Error: '%s' is needed in %s" % (key, func))

            if choices and kwargs[key] not in choices:
                raise Exception("Error: %s :%s" % (key, kwargs[key]))
            ret = func(*args, **kwargs)
            return ret

        return __deco

    return _deco


def dcheck(key, choices=[]):
    def _deco(func):
        def __deco(*args):
            if len(args) == 2:
                values = args[1]
            elif len(args) == 1:
                values = args[0]
            else:
                values = None
            if isinstance(values, dict):
                if key not in values:
                    raise Exception("Error: '%s' is needed in %s" % (key, func))
                if choices and values[key] not in choices:
                    raise Exception("Error: %s :%s" % (key, values[key]))
            ret = func(*args)
            return ret

        return __deco

    return _deco


def vstf_input(key, types=str, choices=[], default=None):
    def _deco(func):
        def __deco(*args):
            ret = func(*args)
            if not isinstance(ret, dict):
                ret = {}
            in_str = "----> %s : " % key
            if choices:
                in_str = "---- %s\n" % (str(choices)) + in_str
            while True:
                if types == list or types == dict:
                    value = input(in_str)
                else:
                    value = raw_input(in_str)
                    value = types(value)
                if not value:
                    value = default
                if not choices or value in choices:
                    break
            ret.update({key: value})
            return ret

        return __deco

    return _deco


def namespace():
    def _deco(func):
        def __deco(*args, **kwargs):
            ret = func(*args, **kwargs)
            nspace = kwargs.get("namespace", None)
            if nspace:
                ret = "ip netns exec %(namespace)s " % {"namespace": nspace} + ret
            return ret

        return __deco

    return _deco