55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
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
|
|
|
|
|
|
PY38 = sys.version_info[:2] >= (3, 8)
|
|
if PY38:
|
|
# On Python 3.8, typed_ast was merged back into `ast`
|
|
_ast_py3 = ast
|
|
|
|
|
|
FunctionType = namedtuple("FunctionType", ["argtypes", "returns"])
|
|
|
|
|
|
def _get_parser_module(parse_python_two=False, type_comments_support=True):
|
|
if not type_comments_support:
|
|
return ast
|
|
|
|
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=False, type_comments=True):
|
|
parse_module = _get_parser_module(
|
|
parse_python_two=parse_python_two, type_comments_support=type_comments
|
|
)
|
|
parse_func = parse_module.parse
|
|
if parse_module is _ast_py3:
|
|
if PY38:
|
|
parse_func = partial(parse_func, type_comments=type_comments)
|
|
if 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)
|