diff --git a/mypy/fastparse.py b/mypy/fastparse.py index f68d1b0279d0..e09e77e59574 100644 --- a/mypy/fastparse.py +++ b/mypy/fastparse.py @@ -22,6 +22,7 @@ ) from mypy import defaults from mypy import experiments +from mypy import hooks from mypy.errors import Errors try: @@ -87,6 +88,33 @@ def parse_type_comment(type_comment: str, line: int) -> Type: return TypeConverter(line=line).visit(typ.body) +def parse_docstring(docstring: str, arg_names: List[str], + line: int) -> Optional[Tuple[List[Type], Type]]: + """Parse a docstring and return type representations. + + Returns a 2-tuple: (list of arguments Types, and return Type). + """ + def pop_and_convert(name): + t = type_map.pop(name, None) + if t is None: + return AnyType() + else: + return parse_type_comment(t, line=line) + + if hooks.docstring_parser is not None: + type_map = hooks.docstring_parser(docstring) + if type_map: + arg_types = [pop_and_convert(name) for name in arg_names] + return_type = pop_and_convert('return') + if type_map: + raise TypeCommentParseError( + 'Arguments parsed from docstring are not present in ' + 'function signature: {}'.format(', '.join(type_map)), + line, 0) + return arg_types, return_type + return None + + def with_line(f: Callable[['ASTConverter', T], U]) -> Callable[['ASTConverter', T], U]: @wraps(f) def wrapper(self: 'ASTConverter', ast: T) -> U: @@ -287,6 +315,14 @@ def do_func_def(self, n: Union[ast35.FunctionDef, ast35.AsyncFunctionDef], else: arg_types = [a.type_annotation for a in args] return_type = TypeConverter(line=n.lineno).visit(n.returns) + # hooks + if (not any(arg_types) and return_type is None and + hooks.docstring_parser): + doc = ast35.get_docstring(n, clean=False) + if doc: + types = parse_docstring(doc, arg_names, n.lineno) + if types is not None: + arg_types, return_type = types for arg, arg_type in zip(args, arg_types): self.set_type_optional(arg_type, arg.initializer) diff --git a/mypy/fastparse2.py b/mypy/fastparse2.py index 76e95bb4cd55..0175d50645fe 100644 --- a/mypy/fastparse2.py +++ b/mypy/fastparse2.py @@ -38,7 +38,9 @@ from mypy import defaults from mypy import experiments from mypy.errors import Errors -from mypy.fastparse import TypeConverter, TypeCommentParseError +from mypy.fastparse import (TypeConverter, TypeCommentParseError, + parse_docstring) +from mypy import hooks try: from typed_ast import ast27 @@ -293,6 +295,15 @@ def visit_FunctionDef(self, n: ast27.FunctionDef) -> Statement: else: arg_types = [a.type_annotation for a in args] return_type = converter.visit(None) + # hooks + if (not any(arg_types) and return_type is None and + hooks.docstring_parser): + doc = ast27.get_docstring(n, clean=False) + if doc: + types = parse_docstring(doc.decode('unicode_escape'), + arg_names, n.lineno) + if types is not None: + arg_types, return_type = types for arg, arg_type in zip(args, arg_types): self.set_type_optional(arg_type, arg.initializer) diff --git a/mypy/hooks.py b/mypy/hooks.py new file mode 100644 index 000000000000..8566c87b7e7d --- /dev/null +++ b/mypy/hooks.py @@ -0,0 +1,13 @@ +from typing import Dict, Optional, Callable + +# The docstring_parser hook is called for each function that has a docstring +# and no other type annotations applied, and the callable should accept the +# docstring as an argument and return a mapping of argument name to type. +# +# The function's return type, if specified, is stored in the mapping with the +# special key 'return'. Other than 'return', the keys of the mapping must be +# a subset of the arguments of the function to which the docstring belongs; an +# error will be raised if the mapping contains stray arguments. +# +# The values of the mapping must be valid PEP484-compatible strings. +docstring_parser = None # type: Callable[[str], Optional[Dict[str, str]]]