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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
import ast
import sys
from collections import namedtuple
from functools import partial
from typing import Dict, Optional
from astroid.const import PY38_PLUS, Context
try:
import typed_ast.ast3 as _ast_py3
except ImportError:
_ast_py3 = None
if PY38_PLUS:
# On Python 3.8, typed_ast was merged back into `ast`
_ast_py3 = ast
FunctionType = namedtuple("FunctionType", ["argtypes", "returns"])
class ParserModule(
namedtuple(
"ParserModule",
[
"module",
"unary_op_classes",
"cmp_op_classes",
"bool_op_classes",
"bin_op_classes",
"context_classes",
],
)
):
def parse(self, string: str, type_comments=True):
if self.module is _ast_py3:
if PY38_PLUS:
parse_func = partial(self.module.parse, type_comments=type_comments)
else:
parse_func = partial(
self.module.parse, feature_version=sys.version_info.minor
)
else:
parse_func = self.module.parse
return parse_func(string)
def parse_function_type_comment(type_comment: str) -> Optional[FunctionType]:
"""Given a correct type comment, obtain a FunctionType object"""
if _ast_py3 is None:
return None
func_type = _ast_py3.parse(type_comment, "<type_comment>", "func_type")
return FunctionType(argtypes=func_type.argtypes, returns=func_type.returns)
def get_parser_module(type_comments=True) -> ParserModule:
if not type_comments:
parser_module = ast
else:
parser_module = _ast_py3
parser_module = parser_module or ast
unary_op_classes = _unary_operators_from_module(parser_module)
cmp_op_classes = _compare_operators_from_module(parser_module)
bool_op_classes = _bool_operators_from_module(parser_module)
bin_op_classes = _binary_operators_from_module(parser_module)
context_classes = _contexts_from_module(parser_module)
return ParserModule(
parser_module,
unary_op_classes,
cmp_op_classes,
bool_op_classes,
bin_op_classes,
context_classes,
)
def _unary_operators_from_module(module):
return {module.UAdd: "+", module.USub: "-", module.Not: "not", module.Invert: "~"}
def _binary_operators_from_module(module):
binary_operators = {
module.Add: "+",
module.BitAnd: "&",
module.BitOr: "|",
module.BitXor: "^",
module.Div: "/",
module.FloorDiv: "//",
module.MatMult: "@",
module.Mod: "%",
module.Mult: "*",
module.Pow: "**",
module.Sub: "-",
module.LShift: "<<",
module.RShift: ">>",
}
return binary_operators
def _bool_operators_from_module(module):
return {module.And: "and", module.Or: "or"}
def _compare_operators_from_module(module):
return {
module.Eq: "==",
module.Gt: ">",
module.GtE: ">=",
module.In: "in",
module.Is: "is",
module.IsNot: "is not",
module.Lt: "<",
module.LtE: "<=",
module.NotEq: "!=",
module.NotIn: "not in",
}
def _contexts_from_module(module) -> Dict[ast.expr_context, Context]:
return {
module.Load: Context.Load,
module.Store: Context.Store,
module.Del: Context.Del,
module.Param: Context.Store,
}
|
