blob: da95f1de354113edf498f7b9fac1208cd1a8d5c3 (
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
|
import ast
from collections import namedtuple
from functools import partial
from typing import Optional
import sys
_ast_py2 = _ast_py3 = None
try:
import typed_ast.ast3 as _ast_py3
import typed_ast.ast27 as _ast_py2
except ImportError:
pass
FunctionType = namedtuple("FunctionType", ["argtypes", "returns"])
def _get_parser_module(parse_python_two: bool = False):
if parse_python_two:
parser_module = _ast_py2
else:
parser_module = _ast_py3
return parser_module or ast
def _parse(string: str, parse_python_two: bool = False):
parse_module = _get_parser_module(parse_python_two=parse_python_two)
parse_func = parse_module.parse
if _ast_py3 and not parse_python_two:
parse_func = partial(parse_func, feature_version=sys.version_info.minor)
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)
|
