Rough draft of pluggable system for producing types from docstrings

python · chadrik · Oct 11, 2016 · Oct 14, 2016 · Oct 15, 2016 · Oct 19, 2016
commit 0a47ddbd5130517ec726b6c4812d8f1bb16f0a95
diff --git a/mypy/docstrings.py b/mypy/docstrings.py
@@ -0,0 +1,27 @@
+from typing import List, Tuple, Dict
+from mypy.types import Type, CallableType, AnyType
+from mypy.nodes import Argument
+
+
+def parse_docstring(docstring: str, line: int) -> Tuple[Dict[str, Type], Type]:
+    """
+    Parse a docstring and return type representations.  This function can
+    be overridden by third-party tools which aim to add typing via docstrings.
+
+    Returns a 2-tuple: dictionary of arg name to Type, and return Type.
+    """
+    return None, None
+
+
+def make_callable(args: List[Argument], type_map: Dict[str, Type],
+                  ret_type: Type) -> CallableType:
+    if type_map is not None:
+        arg_kinds = [arg.kind for arg in args]
+        arg_names = [arg.variable.name() for arg in args]
+        arg_types = [type_map.get(name) for name in arg_names]
+
+        return CallableType([a if a is not None else AnyType() for a in arg_types],
+                            arg_kinds,
+                            arg_names,
+                            ret_type, None,
+                            is_ellipsis_args=False)
diff --git a/mypy/fastparse.py b/mypy/fastparse.py
@@ -1,4 +1,5 @@
 from functools import wraps
+from inspect import cleandoc
 import sys
 
 from typing import Tuple, Union, TypeVar, Callable, Sequence, Optional, Any, cast, List
@@ -22,6 +23,7 @@
 )
 from mypy import defaults
 from mypy import experiments
+from mypy import docstrings
 from mypy.errors import Errors
 
 try:
@@ -287,7 +289,15 @@ 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)
-
+            # docstrings
+            if not any(arg_types) and return_type is None:
+                doc = ast35.get_docstring(n, clean=False)
+                if doc:
+                    doc = cleandoc(doc.decode('unicode_escape'))
+                    type_map, rtype = docstrings.parse_docstring(doc, n.lineno)
+                    if type_map is not None:
+                        arg_types = [type_map.get(name) for name in arg_names]
+                        return_type = rtype
         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
@@ -15,6 +15,7 @@
 two in a typesafe way.
 """
 from functools import wraps
+from inspect import cleandoc
 import sys
 
 from typing import Tuple, Union, TypeVar, Callable, Sequence, Optional, Any, cast, List
@@ -39,6 +40,7 @@
 from mypy import experiments
 from mypy.errors import Errors
 from mypy.fastparse import TypeConverter, TypeCommentParseError
+from mypy import docstrings
 
 try:
     from typed_ast import ast27
@@ -293,7 +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)
-
+            # docstrings
+            if not any(arg_types) and return_type is None:
+                doc = ast27.get_docstring(n, clean=False)
+                if doc:
+                    doc = cleandoc(doc.decode('unicode_escape'))
+                    type_map, rtype = docstrings.parse_docstring(doc, n.lineno)
+                    if type_map is not None:
+                        arg_types = [type_map.get(name) for name in arg_names]
+                        return_type = rtype
         for arg, arg_type in zip(args, arg_types):
             self.set_type_optional(arg_type, arg.initializer)
 

diff --git a/mypy/parse.py b/mypy/parse.py
@@ -5,6 +5,7 @@
 """
 
 import re
+from inspect import cleandoc
 
 from typing import List, Tuple, Set, cast, Union, Optional
 
@@ -38,6 +39,7 @@
 from mypy.options import Options
 
 from mypy import experiments
+from mypy import docstrings
 
 
 class ParseError(Exception): pass
@@ -417,7 +419,7 @@ def parse_function(self, no_type_checks: bool = False) -> FuncDef:
             arg_kinds = [arg.kind for arg in args]
             arg_names = [arg.variable.name() for arg in args]
 
-            body, comment_type = self.parse_block(allow_type=True)
+            body, comment_type = self.parse_block(args=args)
             # Potentially insert extra assignment statements to the beginning of the
             # body, used to decompose Python 2 tuple arguments.
             body.body[:0] = extra_stmts
@@ -831,7 +833,7 @@ def construct_function_type(self, args: List[Argument], ret_type: Type,
 
     # Parsing statements
 
-    def parse_block(self, allow_type: bool = False) -> Tuple[Block, Type]:
+    def parse_block(self, args: List[Argument]=None) -> Tuple[Block, Type]:
         colon = self.expect(':')
         if not isinstance(self.current(), Break):
             # Block immediately after ':'.
@@ -854,6 +856,13 @@ def parse_block(self, allow_type: bool = False) -> Tuple[Block, Type]:
             brk = self.expect_break()
             type = self.parse_type_comment(brk, signature=True)
             self.expect_indent()
+            if args is not None:
+                cur = self.current()
+                if type is None and isinstance(cur, StrLit):
+                    type_map, ret_type = docstrings.parse_docstring(
+                        cleandoc(cur.parsed()), cur.line)
+                    type = docstrings.make_callable(args, type_map, ret_type)
+
             stmt_list = []  # type: List[Statement]
             while (not isinstance(self.current(), Dedent) and
                    not isinstance(self.current(), Eof)):