bedrock.generators.introspection.decorators

 1import inspect
 2
 3
 4class Decorator:
 5    def __init__(self, decorator_string):
 6        self.string = decorator_string
 7        self.name = decorator_string.split("(")[0]
 8        self.kwargs = {}
 9        self.args = []
10        if "(" in decorator_string:
11            all_args = decorator_string.split("(")[1].split(")")[0].split(",")
12            for arg in all_args:
13                if "=" in arg:
14                    key = arg.split("=")[0].strip()
15                    value = arg.split("=")[1].strip()
16                    self.kwargs[_resolve_decorator_arg(key)] = _resolve_decorator_arg(value)
17                else:
18                    self.args.append(_resolve_decorator_arg(arg))
19
20    def __repr__(self):
21        return self.string
22
23
24def get_function_decorators(function) -> list[Decorator]:
25    """
26    Get decorators for the given function.
27    :param function: Function to get decorators from.
28    :return: List of decorator names.
29    """
30    source = inspect.getsource(function)
31    index = source.find("def ")
32    decorator_line = [line.strip() for line in source[:index].strip().splitlines() if line.strip()[0] == "@"]
33    return [Decorator(decorator) for decorator in decorator_line]
34
35
36def get_class_decorators(cls) -> list[Decorator]:
37    """
38    Get decorators for the given Class.
39    :param cls: Class to get decorators from.
40    :return: List of decorator names.
41    """
42    source = inspect.getsource(cls)
43    index = source.find("class ")
44    decorator_line = [line.strip() for line in source[:index].strip().splitlines() if line.strip()[0] == "@"]
45    return [Decorator(decorator) for decorator in decorator_line]
46
47
48def _resolve_decorator_arg(arg):
49    _arg = arg.strip()
50    types = {
51        "str": str,
52        "integer": int,
53        "int": int,
54        "float": float,
55        "bool": bool,
56        "list": list,
57        "dict": dict
58    }
59    if "'" in _arg or '"' in _arg:
60        return _arg.strip("'").strip('"')
61    elif _arg in types:
62        return types[_arg]
63    else:
64        try:
65            return int(_arg)
66        except ValueError:
67            try:
68                return float(_arg)
69            except ValueError:
70                return _arg