Pluggable system for producing types from docstrings

Author: chadrik

mypy/fastparse.py

@@ -27,6 +27,7 @@
 from mypy import defaults
 from mypy import experiments
 from mypy import messages
+from mypy import hooks
 from mypy.errors import Errors
 
 try:
@@ -55,6 +56,8 @@
 
 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,
@@ -109,6 +112,33 @@ def parse_type_comment(type_comment: str, line: int, errors: Errors) -> Optional
         return TypeConverter(errors, line=line).visit(typ.body)
 
 
+def parse_docstring(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).
+    """
+    opts = hooks.options.get('docstring_parser', {})
+
+    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], line + t[1], errors)
+
+    if hooks.docstring_parser is not None:
+        type_map = hooks.docstring_parser(docstring, opts, 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(type_map.keys(), 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:
@@ -301,8 +331,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]
@@ -342,6 +373,14 @@ def do_func_def(self, n: Union[ast3.FunctionDef, ast3.AsyncFunctionDef],
         else:
             arg_types = [a.type_annotation for a in args]
             return_type = TypeConverter(self.errors, line=n.lineno).visit(n.returns)
+            # hooks
+            if (not any(arg_types) and return_type is None and
+                    hooks.docstring_parser):
+                doc = ast3.get_docstring(n, clean=False)
+                if doc:
+                    types = parse_docstring(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

@@ -42,7 +42,9 @@
 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 import hooks
 
 try:
     from typed_ast import ast27
@@ -290,8 +292,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)
 
@@ -326,6 +329,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'),
+                                            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/hooks.py

@@ -0,0 +1,22 @@
+from typing import Dict, Optional, Callable, Tuple
+from mypy.errors import Errors
+
+Options = Dict[str, str]
+options = {}  # type: Dict[str, Options]
+
+# The docstring_parser hook is called for each unannotated function that has a
+# docstring. The callable should accept three arguments:
+# - the docstring to be parsed
+# - a dictionary of options (parsed from the [docstring_parser] section of mypy
+#   config file)
+# - an Errors object for reporting errors, warnings, and info.
+#
+# The function should return a map from argument name to 2-tuple.  The latter should contain:
+# - a PEP484-compatible string. 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.
+# - a line number offset, relative to the start of the docstring, used to
+#   improve errors if the associated type string is invalid.
+#
+docstring_parser = None  # type: Callable[[str, Options, Errors], Optional[Dict[str, Tuple[str, int]]]]

mypy/main.py

@@ -7,12 +7,14 @@
 import re
 import sys
 import time
+from pydoc import locate
 
-from typing import Any, Dict, List, Mapping, Optional, Set, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Set, Tuple, Callable
 
 from mypy import build
 from mypy import defaults
 from mypy import experiments
+from mypy import hooks
 from mypy import util
 from mypy.build import BuildSource, BuildResult, PYTHON_EXTENSIONS
 from mypy.errors import CompileError
@@ -548,6 +550,19 @@ def get_init_file(dir: str) -> Optional[str]:
     return None
 
 
+def load_hook(prefix: str, hook_name: str, hook_path: str) -> Optional[Callable]:
+    # FIXME: no stubs for pydoc.  should we write stubs or a simple replacement for locate?
+    obj = locate(hook_path)
+    if obj is None:
+        print("%s: Could not find hook %s at %s" %
+              (prefix, hook_name, hook_path), file=sys.stderr)
+    if not callable(obj):
+        print("%s: Hook %s at %s is not callable" %
+              (prefix, hook_name, hook_path), file=sys.stderr)
+        return None
+    return obj
+
+
 # For most options, the type of the default value set in options.py is
 # sufficient, and we don't have to do anything here.  This table
 # exists to specify types for values initialized to None or container
@@ -603,19 +618,33 @@ def parse_config_file(options: Options, filename: Optional[str]) -> None:
     else:
         section = parser['mypy']
         prefix = '%s: [%s]' % (file_read, 'mypy')
-        updates, report_dirs = parse_section(prefix, options, section)
+        updates, report_dirs, hook_funcs = parse_section(prefix, options, section)
         for k, v in updates.items():
             setattr(options, k, v)
+
+        # bind hook functions to hooks module
+        for k, v in hook_funcs.items():
+            hook_func = load_hook(prefix, k, v)
+            if hook_func is not None:
+                # FIXME: dynamically check loaded function annotations against those in `hooks`?
+                setattr(hooks, k, hook_func)
+                # look for an options section for this hook
+                if k in parser:
+                    hooks.options[k] = dict(parser[k])
         options.report_dirs.update(report_dirs)
 
     for name, section in parser.items():
         if name.startswith('mypy-'):
             prefix = '%s: [%s]' % (file_read, name)
-            updates, report_dirs = parse_section(prefix, options, section)
+            updates, report_dirs, hook_funcs = parse_section(prefix, options, section)
             if report_dirs:
                 print("%s: Per-module sections should not specify reports (%s)" %
                       (prefix, ', '.join(s + '_report' for s in sorted(report_dirs))),
                       file=sys.stderr)
+            if hook_funcs:
+                print("%s: Per-module sections should not specify hooks (%s)" %
+                      (prefix, ', '.join(sorted(hook_funcs))),
+                      file=sys.stderr)
             if set(updates) - Options.PER_MODULE_OPTIONS:
                 print("%s: Per-module sections should only specify per-module flags (%s)" %
                       (prefix, ', '.join(sorted(set(updates) - Options.PER_MODULE_OPTIONS))),
@@ -632,16 +661,27 @@ def parse_config_file(options: Options, filename: Optional[str]) -> None:
 
 
 def parse_section(prefix: str, template: Options,
-                  section: Mapping[str, str]) -> Tuple[Dict[str, object], Dict[str, str]]:
+                  section: Mapping[str, str]) -> Tuple[Dict[str, object],
+                                                       Dict[str, str], Dict[str, str]]:
     """Parse one section of a config file.
 
     Returns a dict of option values encountered, and a dict of report directories.
     """
     results = {}  # type: Dict[str, object]
     report_dirs = {}  # type: Dict[str, str]
+    hook_funcs = {}  # type: Dict[str, str]
     for key in section:
         key = key.replace('-', '_')
-        if key in config_types:
+        if key.startswith('hooks.'):
+            dv = section.get(key)
+            key = key[6:]
+            if not hasattr(hooks, key):
+                print("%s: Unrecognized hook: %s = %s" % (prefix, key, dv),
+                      file=sys.stderr)
+            else:
+                hook_funcs[key] = dv
+            continue
+        elif key in config_types:
             ct = config_types[key]
         else:
             dv = getattr(template, key, None)
@@ -685,7 +725,7 @@ def parse_section(prefix: str, template: Options,
                 if 'follow_imports' not in results:
                     results['follow_imports'] = 'error'
         results[key] = v
-    return results, report_dirs
+    return results, report_dirs, hook_funcs
 
 
 def fail(msg: str) -> None: