Completed
Push — master ( 2020be...46184f )
by Max
02:41
created

structured_data._ctor.annotation_is_classvar()   A

Complexity

Conditions 3

Size

Total Lines 7
Code Lines 7

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 7
nop 2
dl 0
loc 7
rs 10
c 0
b 0
f 0
1
"""Internal implementation of helper types for Sum annotations."""
2
3
import ast
4
import typing
5
import weakref
6
7
import astor  # type: ignore
8
9
_CTOR_CACHE: typing.Dict[typing.Tuple, "Ctor"] = {}
10
11
12
ARGS: typing.MutableMapping[
13
    typing.Union["Ctor", typing.Type["Ctor"]], typing.Tuple
14
] = weakref.WeakKeyDictionary()
15
16
17
class Ctor:
18
    """Marker class for adt constructors.
19
20
    To use, index with a sequence of types, and annotate a variable in an
21
    adt-decorated class with it.
22
    """
23
24
    __slots__ = ("__weakref__",)
25
26
    def __new__(cls, args):
27
        if args == ():
28
            return cls
29
        self = object.__new__(cls)
30
        ARGS[self] = args
31
        return _CTOR_CACHE.setdefault(args, self)
32
33
    def __init_subclass__(cls, **kwargs):
34
        raise TypeError
35
36
    def __class_getitem__(cls, args):
37
        if not isinstance(args, tuple):
38
            args = (args,)
39
        # Yes it is.
40
        return cls(args)  # pylint: disable=not-callable
41
42
43
ARGS[Ctor] = ()
44
45
46
def _interpret_args_from_non_string(
47
    constructor: typing.Any
48
) -> typing.Optional[typing.Tuple]:
49
    try:
50
        return ARGS.get(constructor)
51
    except TypeError:
52
        return None
53
54
55
def _interpret_classvar_from_non_string(annotation: typing.Any) -> bool:
56
    if annotation is typing.ClassVar:
57
        return True
58
    try:
59
        return annotation.__origin__ is typing.ClassVar
60
    except AttributeError:
61
        return False
62
63
64
def _parse_constructor(constructor: str) -> ast.Module:
65
    try:
66
        return ast.parse(constructor, mode="eval")
67
    except Exception:
68
        raise ValueError("parsing annotation failed")
69
70
71
def _get_args_from_index(index: ast.AST) -> typing.Tuple:
72
    if isinstance(index, ast.Tuple):
73
        return tuple(astor.to_source(elt) for elt in index.elts)
74
    return (astor.to_source(index),)
75
76
77
def _checked_eval(source, global_ns: typing.Dict[str, typing.Any]) -> typing.Any:
78
    try:
79
        # Oh no, the user might end up executing arbitrary code that they wrote
80
        # in the first place.
81
        return eval(source, global_ns)  # pylint: disable=eval-used
82
    # If we hit this, anything could have gone wrong, but just assume it's not
83
    # anything we care about.
84
    except Exception:  # pylint: disable=broad-except
85
        return None
86
87
88
NO_VALUE = object()
89
90
91
def _extract_tuple_ast(
92
    constructor: str, global_ns: typing.Dict[str, typing.Any]
93
) -> typing.Optional[typing.Tuple]:
94
    ctor_ast = _parse_constructor(constructor)
95
    value = index = NO_VALUE
96
    if isinstance(ctor_ast.body, ast.Subscript) and isinstance(
97
        ctor_ast.body.slice, ast.Index
98
    ):
99
        index = ctor_ast.body.slice.value
100
        ctor_ast.body = ctor_ast.body.value
101
        value = _checked_eval(compile(ctor_ast, "<annotation>", "eval"), global_ns)
102
    if value is Ctor:
103
        # If value is Ctor, then value was set, which is only possible if the
104
        # previous block executed, which will set "index" to an AST.
105
        return _get_args_from_index(typing.cast(ast.AST, index))
106
    if value is None:
107
        return None
108
    return _interpret_args_from_non_string(_checked_eval(constructor, global_ns))
109
110
111
def _str_is_classvar(annotation: str, global_ns: typing.Dict[str, typing.Any]) -> bool:
112
    annotation_ast = _parse_constructor(annotation)
113
    value = NO_VALUE
114
    if isinstance(annotation_ast.body, ast.Subscript) and isinstance(
115
        annotation_ast.body.slice, ast.Index
116
    ):
117
        annotation_ast.body = annotation_ast.body.value
118
        value = _checked_eval(
119
            compile(annotation_ast, "<annotation>", "eval"), global_ns
120
        )
121
    if value is typing.ClassVar:
122
        return True
123
    if value is None:
124
        return False
125
    return _interpret_classvar_from_non_string(_checked_eval(annotation, global_ns))
126
127
128
def get_args(
129
    constructor, global_ns: typing.Dict[str, typing.Any]
130
) -> typing.Optional[typing.Tuple]:
131
    """Given annotation value and module namespace, return Ctor args, if any.
132
133
    The function first checks if the value is a string. If so, it tries to
134
    parse it.
135
    Otherwise, if the value is a Ctor instance, it extracts the annotation, and
136
    if not, it returns ``None``.
137
    """
138
    if isinstance(constructor, str):
139
        try:
140
            return _extract_tuple_ast(constructor, global_ns)
141
        except ValueError:
142
            return None
143
    return _interpret_args_from_non_string(constructor)
144
145
146
def annotation_is_classvar(annotation, global_ns: typing.Dict[str, typing.Any]) -> bool:
147
    if isinstance(annotation, str):
148
        try:
149
            return _str_is_classvar(annotation, global_ns)
150
        except ValueError:
151
            return False
152
    return _interpret_classvar_from_non_string(annotation)
153