Add plugin hook for extracting types from docstrings

Author: chadrik

mypy/build.py

@@ -532,7 +532,7 @@ def parse_file(self, id: str, path: str, source: str, ignore_errors: bool) -> My
         Raise CompileError if there is a parse error.
         """
         num_errs = self.errors.num_messages()
-        tree = parse(source, path, self.errors, options=self.options)
+        tree = parse(source, path, self.errors, options=self.options, plugin=self.plugin)
         tree._fullname = id
 
         if self.errors.num_messages() != num_errs:

mypy/fastparse.py

@@ -31,6 +31,7 @@
 from mypy import messages
 from mypy.errors import Errors
 from mypy.options import Options
+from mypy.plugin import Plugin, DocstringParserHook
 
 try:
     from typed_ast import ast3
@@ -58,10 +59,12 @@
 
 TYPE_COMMENT_SYNTAX_ERROR = 'syntax error in type comment'
 TYPE_COMMENT_AST_ERROR = 'invalid type comment or annotation'
+TYPE_COMMENT_DOCSTRING_ERROR = ('Arguments parsed from docstring are not '
+                                'present in function signature: {} not in {}')
 
 
 def parse(source: Union[str, bytes], fnam: str = None, errors: Errors = None,
-          options: Options = Options()) -> MypyFile:
+          options: Options = Options(), plugin: Plugin = Plugin()) -> MypyFile:
 
     """Parse a source file, without doing any semantic analysis.
 
@@ -85,6 +88,7 @@ def parse(source: Union[str, bytes], fnam: str = None, errors: Errors = None,
         tree = ASTConverter(options=options,
                             is_stub=is_stub_file,
                             errors=errors,
+                            plugin=plugin,
                             ).visit(ast)
         tree.path = fnam
         tree.is_stub = is_stub_file
@@ -112,6 +116,31 @@ def parse_type_comment(type_comment: str, line: int, errors: Optional[Errors]) -
         return TypeConverter(errors, line=line).visit(typ.body)
 
 
+def parse_docstring(hook: DocstringParserHook, docstring: str, arg_names: List[str],
+                    line: int, errors: Errors) -> 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: str) -> Optional[Type]:
+        t = type_map.pop(name, None)
+        if t is None:
+            return AnyType()
+        else:
+            return parse_type_comment(t[0], t[1], errors)
+
+    type_map = hook(docstring, line, errors)
+    if type_map:
+        arg_types = [pop_and_convert(name) for name in arg_names]
+        return_type = pop_and_convert('return')
+        if type_map:
+            errors.report(line, 0,
+                          TYPE_COMMENT_DOCSTRING_ERROR.format(list(type_map), arg_names))
+        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:
@@ -141,13 +170,15 @@ class ASTConverter(ast3.NodeTransformer):  # type: ignore  # typeshed PR #931
     def __init__(self,
                  options: Options,
                  is_stub: bool,
-                 errors: Errors) -> None:
+                 errors: Errors,
+                 plugin: Plugin) -> None:
         self.class_nesting = 0
         self.imports = []  # type: List[ImportBase]
 
         self.options = options
         self.is_stub = is_stub
         self.errors = errors
+        self.plugin = plugin
 
     def fail(self, msg: str, line: int, column: int) -> None:
         self.errors.report(line, column, msg)
@@ -302,8 +333,9 @@ def do_func_def(self, n: Union[ast3.FunctionDef, ast3.AsyncFunctionDef],
         args = self.transform_args(n.args, n.lineno, no_type_check=no_type_check)
 
         arg_kinds = [arg.kind for arg in args]
-        arg_names = [arg.variable.name() for arg in args]  # type: List[Optional[str]]
-        arg_names = [None if argument_elide_name(name) else name for name in arg_names]
+        real_names = [arg.variable.name() for arg in args]  # type: List[str]
+        arg_names = [None if argument_elide_name(name) else name
+                     for name in real_names]  # type: List[Optional[str]]
         if special_function_elide_names(n.name):
             arg_names = [None] * len(arg_names)
         arg_types = None  # type: List[Type]
@@ -345,6 +377,14 @@ def do_func_def(self, n: Union[ast3.FunctionDef, ast3.AsyncFunctionDef],
             return_type = TypeConverter(self.errors, line=n.returns.lineno
                                         if n.returns else n.lineno).visit(n.returns)
 
+            docstring_hook = self.plugin.get_docstring_parser_hook()
+            if docstring_hook is not None and not any(arg_types) and return_type is None:
+                doc = ast3.get_docstring(n, clean=False)
+                if doc:
+                    types = parse_docstring(docstring_hook, doc, real_names, n.lineno, self.errors)
+                    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)
 

mypy/fastparse2.py

@@ -41,8 +41,10 @@
 from mypy import experiments
 from mypy import messages
 from mypy.errors import Errors
-from mypy.fastparse import TypeConverter, parse_type_comment
+from mypy.fastparse import (TypeConverter, parse_type_comment,
+                            parse_docstring)
 from mypy.options import Options
+from mypy.plugin import Plugin
 
 try:
     from typed_ast import ast27
@@ -74,7 +76,7 @@
 
 
 def parse(source: Union[str, bytes], fnam: str = None, errors: Errors = None,
-          options: Options = Options()) -> MypyFile:
+          options: Options = Options(), plugin: Plugin = Plugin()) -> MypyFile:
     """Parse a source file, without doing any semantic analysis.
 
     Return the parse tree. If errors is not provided, raise ParseError
@@ -92,6 +94,7 @@ def parse(source: Union[str, bytes], fnam: str = None, errors: Errors = None,
         tree = ASTConverter(options=options,
                             is_stub=is_stub_file,
                             errors=errors,
+                            plugin=plugin,
                             ).visit(ast)
         assert isinstance(tree, MypyFile)
         tree.path = fnam
@@ -135,13 +138,15 @@ class ASTConverter(ast27.NodeTransformer):
     def __init__(self,
                  options: Options,
                  is_stub: bool,
-                 errors: Errors) -> None:
+                 errors: Errors,
+                 plugin: Plugin) -> None:
         self.class_nesting = 0
         self.imports = []  # type: List[ImportBase]
 
         self.options = options
         self.is_stub = is_stub
         self.errors = errors
+        self.plugin = plugin
 
     def fail(self, msg: str, line: int, column: int) -> None:
         self.errors.report(line, column, msg)
@@ -284,8 +289,9 @@ def visit_FunctionDef(self, n: ast27.FunctionDef) -> Statement:
         args, decompose_stmts = self.transform_args(n.args, n.lineno)
 
         arg_kinds = [arg.kind for arg in args]
-        arg_names = [arg.variable.name() for arg in args]  # type: List[Optional[str]]
-        arg_names = [None if argument_elide_name(name) else name for name in arg_names]
+        real_names = [arg.variable.name() for arg in args]  # type: List[str]
+        arg_names = [None if argument_elide_name(name) else name
+                     for name in real_names]  # type: List[Optional[str]]
         if special_function_elide_names(n.name):
             arg_names = [None] * len(arg_names)
 
@@ -321,6 +327,15 @@ def visit_FunctionDef(self, n: ast27.FunctionDef) -> Statement:
             arg_types = [a.type_annotation for a in args]
             return_type = converter.visit(None)
 
+            docstring_hook = self.plugin.get_docstring_parser_hook()
+            if docstring_hook is not None and not any(arg_types) and return_type is None:
+                doc = ast27.get_docstring(n, clean=False)
+                if doc:
+                    types = parse_docstring(docstring_hook, doc.decode('unicode_escape'),
+                                            real_names, n.lineno, self.errors)
+                    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)
 

mypy/parse.py

@@ -2,13 +2,15 @@
 
 from mypy.errors import Errors
 from mypy.options import Options
+from mypy.plugin import Plugin
 from mypy.nodes import MypyFile
 
 
 def parse(source: Union[str, bytes],
           fnam: str,
           errors: Optional[Errors],
-          options: Options) -> MypyFile:
+          options: Options,
+          plugin: Plugin) -> MypyFile:
     """Parse a source file, without doing any semantic analysis.
 
     Return the parse tree. If errors is not provided, raise ParseError
@@ -22,10 +24,12 @@ def parse(source: Union[str, bytes],
         return mypy.fastparse.parse(source,
                                     fnam=fnam,
                                     errors=errors,
-                                    options=options)
+                                    options=options,
+                                    plugin=plugin)
     else:
         import mypy.fastparse2
         return mypy.fastparse2.parse(source,
                                      fnam=fnam,
                                      errors=errors,
-                                     options=options)
+                                     options=options,
+                                     plugin=plugin)

mypy/plugin.py

@@ -1,12 +1,16 @@
-from typing import Callable, List, Tuple, Optional, NamedTuple, TypeVar
+from typing import Callable, Dict, List, Tuple, Optional, NamedTuple, TypeVar
 
 from mypy.nodes import Expression, StrExpr, IntExpr, UnaryExpr, Context
 from mypy.types import (
     Type, Instance, CallableType, TypedDictType, UnionType, NoneTyp, FunctionLike, TypeVarType,
     AnyType
 )
 from mypy.messages import MessageBuilder
+from mypy import defaults
 
+MYPY = False
+if MYPY:
+    from mypy.errors import Errors
 
 # Create an Instance given full name of class and type arguments.
 NamedInstanceCallback = Callable[[str, List[Type]], Type]
@@ -58,6 +62,27 @@
     Type  # Return type inferred by the callback
 ]
 
+# A callback for extracting type annotations from docstrings.
+#
+# Called for each unannotated function that has adocstring.
+# The function's return type, if specified, is stored in the mapping with the special
+# key 'return'.  Other than 'return', each key of the mapping must be one of the
+# arguments of the documented function; otherwise, an error will be raised.
+DocstringParserHook = Callable[
+    [
+        str,                     # The docstring to be parsed
+        int,                     # The line number where the docstring begins
+        'Errors'                 # Errors object for reporting errors, warnings, and info
+    ],
+    Dict[
+        str,                     # Argument name, or 'return' for return type.
+        Tuple[
+            str,                 # PEP484-compatible string to use as the argument's type
+            int                  # Line number identifying the location of the type annotation
+        ]
+    ]
+]
+
 
 class Plugin:
     """Base class of all type checker plugins.
@@ -69,7 +94,7 @@ class Plugin:
     results might be cached).
     """
 
-    def __init__(self, python_version: Tuple[int, int]) -> None:
+    def __init__(self, python_version: Tuple[int, int] = defaults.PYTHON3_VERSION) -> None:
         self.python_version = python_version
 
     def get_function_hook(self, fullname: str) -> Optional[FunctionHook]:
@@ -81,6 +106,9 @@ def get_method_signature_hook(self, fullname: str) -> Optional[MethodSignatureHo
     def get_method_hook(self, fullname: str) -> Optional[MethodHook]:
         return None
 
+    def get_docstring_parser_hook(self) -> Optional[DocstringParserHook]:
+        return None
+
     # TODO: metaclass / class decorator hook
 
 
@@ -116,6 +144,9 @@ def get_method_signature_hook(self, fullname: str) -> Optional[MethodSignatureHo
     def get_method_hook(self, fullname: str) -> Optional[MethodHook]:
         return self._find_hook(lambda plugin: plugin.get_method_hook(fullname))
 
+    def get_docstring_parser_hook(self) -> Optional[DocstringParserHook]:
+        return self._find_hook(lambda plugin: plugin.get_docstring_parser_hook())
+
     def _find_hook(self, lookup: Callable[[Plugin], T]) -> Optional[T]:
         for plugin in self._plugins:
             hook = lookup(plugin)

mypy/stubgen.py

@@ -63,6 +63,7 @@
 from mypy.stubgenc import parse_all_signatures, find_unique_signatures, generate_stub_for_c_module
 from mypy.stubutil import is_c_module, write_header
 from mypy.options import Options as MypyOptions
+from mypy.plugin import Plugin
 
 
 Options = NamedTuple('Options', [('pyversion', Tuple[int, int]),
@@ -194,8 +195,9 @@ def generate_stub(path: str, output_dir: str, _all_: Optional[List[str]] = None,
         source = f.read()
     options = MypyOptions()
     options.python_version = pyversion
+    plugin = Plugin(pyversion)
     try:
-        ast = mypy.parse.parse(source, fnam=path, errors=None, options=options)
+        ast = mypy.parse.parse(source, fnam=path, errors=None, options=options, plugin=plugin)
     except mypy.errors.CompileError as e:
         # Syntax error!
         for m in e.messages:

mypy/test/testcheck.py

@@ -77,6 +77,7 @@
     'check-enum.test',
     'check-incomplete-fixture.test',
     'check-custom-plugin.test',
+    'check-docstring-hook.test',
 ]
 
 

mypy/test/testparse.py

@@ -12,6 +12,7 @@
 from mypy.parse import parse
 from mypy.errors import CompileError
 from mypy.options import Options
+from mypy.plugin import Plugin
 
 
 class ParserSuite(Suite):
@@ -43,7 +44,8 @@ def test_parser(testcase: DataDrivenTestCase) -> None:
         n = parse(bytes('\n'.join(testcase.input), 'ascii'),
                   fnam='main',
                   errors=None,
-                  options=options)
+                  options=options,
+                  plugin=Plugin(options.python_version))
         a = str(n).split('\n')
     except CompileError as e:
         a = e.messages
@@ -68,7 +70,8 @@ def cases(self) -> List[DataDrivenTestCase]:
 def test_parse_error(testcase: DataDrivenTestCase) -> None:
     try:
         # Compile temporary file. The test file contains non-ASCII characters.
-        parse(bytes('\n'.join(testcase.input), 'utf-8'), INPUT_FILE_NAME, None, Options())
+        parse(bytes('\n'.join(testcase.input), 'utf-8'), INPUT_FILE_NAME, None, Options(),
+              Plugin())
         raise AssertionFailure('No errors reported')
     except CompileError as e:
         # Verify that there was a compile error and that the error messages