Pluggable system for generating types from docstrings

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)

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)

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]]]