create type hint stub files for module torch

.gitignore

@@ -35,11 +35,13 @@ test/data/gpu_tensors.pt
 test/data/legacy_modules.t7
 test/data/legacy_serialized.pt
 test/data/linear.pt
+test/generated_type_hints_smoketest.py
 test/htmlcov
 test/cpp_extensions/install/
 third_party/build/
 tools/shared/_utils_internal.py
 torch.egg-info/
+torch/__init__.pyi
 torch/csrc/autograd/generated/*
 torch/csrc/cudnn/cuDNN.cpp
 torch/csrc/generated

.jenkins/pytorch/test.sh

@@ -34,6 +34,10 @@ if [[ "$BUILD_ENVIRONMENT" != *ppc64le* ]]; then
 
   # TODO: move this to Docker
   pip install -q hypothesis --user
+
+  # mypy will fail to install on Python <3.4.  In that case,
+  # we just won't run these tests.
+  pip install mypy --user || true
 fi
 
 # DANGER WILL ROBINSON.  The LD_PRELOAD here could cause you problems

setup.py

@@ -731,6 +731,7 @@ def print_box(msg):
         entry_points=entry_points,
         package_data={
             'torch': [
+                '__init__.pyi',
                 'lib/*.so*',
                 'lib/*.dylib*',
                 'lib/*.dll',

test/run_test.py

@@ -42,6 +42,7 @@
     'thd_distributed',
     'torch',
     'type_info',
+    'type_hints',
     'utils',
 ]
 

test/test_type_hints.py

@@ -0,0 +1,167 @@
+from __future__ import print_function
+import unittest
+from common_utils import TestCase, run_tests, download_file
+import tempfile
+import torch
+import re
+import os
+import sys
+import subprocess
+import inspect
+
+try:
+    import mypy
+    HAVE_MYPY = True
+except ImportError:
+    HAVE_MYPY = False
+
+
+def get_examples_from_docstring(docstr):
+    """
+    Extracts all runnable python code from the examples
+    in docstrings; returns a list of lines.
+    """
+    # TODO: Figure out if there's a way to use doctest directly to
+    # implement this
+    example_file_lines = []
+    # the detection is a bit hacky because there isn't a nice way of detecting
+    # where multiline commands end. Thus we keep track of how far we got in beginning
+    # and continue to add lines until we have a compileable Python statement.
+    exampleline_re = re.compile(r"^\s+(?:>>>|\.\.\.) (.*)$")
+    beginning = ""
+    for l in docstr.split('\n'):
+        if beginning:
+            m = exampleline_re.match(l)
+            if m:
+                beginning += m.group(1)
+            else:
+                beginning += l
+        else:
+            m = exampleline_re.match(l)
+            if m:
+                beginning += m.group(1)
+        if beginning:
+            complete = True
+            try:
+                compile(beginning, "", "exec")
+            except SyntaxError:
+                complete = False
+            if complete:
+                # found one
+                example_file_lines += beginning.split('\n')
+                beginning = ""
+            else:
+                beginning += "\n"
+    return ['    ' + l for l in example_file_lines]
+
+
+def get_all_examples():
+    """get_all_examples() -> str
+
+    This function grabs (hopefully all) examples from the torch documentation
+    strings and puts them in one nonsensical module returned as a string.
+    """
+    blacklist = {"_np"}
+    allexamples = ""
+
+    example_file_lines = [
+        "import torch",
+        "import torch.nn.functional as F",
+        "import math  # type: ignore",  # mypy complains about floats where SupportFloat is expected
+        "import numpy  # type: ignore",
+        "import io  # type: ignore",
+        "import itertools  # type: ignore",
+        "",
+        # for requires_grad_ example
+        # NB: We are parsing this file as Python 2, so we must use
+        # Python 2 type annotation syntax
+        "def preprocess(inp):",
+        "    # type: (torch.Tensor) -> torch.Tensor",
+        "    return inp",
+    ]
+
+    for fname in dir(torch):
+        fn = getattr(torch, fname)
+        docstr = inspect.getdoc(fn)
+        if docstr and fname not in blacklist:
+            e = get_examples_from_docstring(docstr)
+            if e:
+                example_file_lines.append("\n\ndef example_torch_{}():".format(fname))
+                example_file_lines += e
+
+    for fname in dir(torch.Tensor):
+        fn = getattr(torch.Tensor, fname)
+        docstr = inspect.getdoc(fn)
+        if docstr and fname not in blacklist:
+            e = get_examples_from_docstring(docstr)
+            if e:
+                example_file_lines.append("\n\ndef example_torch_tensor_{}():".format(fname))
+                example_file_lines += e
+
+    return "\n".join(example_file_lines)
+
+
+class TestTypeHints(TestCase):
+    @unittest.skipIf(sys.version_info[0] == 2, "no type hints for Python 2")
+    @unittest.skipIf(not HAVE_MYPY, "need mypy")
+    def test_doc_examples(self):
+        """
+        Run documentation examples through mypy.
+        """
+        fn = os.path.join(os.path.dirname(__file__), 'generated_type_hints_smoketest.py')
+        with open(fn, "w") as f:
+            print(get_all_examples(), file=f)
+
+        # OK, so here's the deal.  mypy treats installed packages
+        # and local modules differently: if a package is installed,
+        # mypy will refuse to use modules from that package for type
+        # checking unless the module explicitly says that it supports
+        # type checking. (Reference:
+        # https://mypy.readthedocs.io/en/latest/running_mypy.html#missing-imports
+        # )
+        #
+        # Now, PyTorch doesn't support typechecking, and we shouldn't
+        # claim that it supports typechecking (it doesn't.) However, not
+        # claiming we support typechecking is bad for this test, which
+        # wants to use the partial information we get from the bits of
+        # PyTorch which are typed to check if it typechecks.  And
+        # although mypy will work directly if you are working in source,
+        # some of our tests involve installing PyTorch and then running
+        # its tests.
+        #
+        # The guidance we got from Michael Sullivan and Joshua Oreman,
+        # and also independently developed by Thomas Viehmann,
+        # is that we should create a fake directory and add symlinks for
+        # the packages that should typecheck.  So that is what we do
+        # here.
+        #
+        # If you want to run mypy by hand, and you run from PyTorch
+        # root directory, it should work fine to skip this step (since
+        # mypy will preferentially pick up the local files first).  The
+        # temporary directory here is purely needed for CI.  For this
+        # reason, we also still drop the generated file in the test
+        # source folder, for ease of inspection when there are failures.
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            try:
+                os.symlink(
+                    os.path.dirname(torch.__file__),
+                    os.path.join(tmp_dir, 'torch'),
+                    target_is_directory=True
+                )
+            except OSError:
+                raise unittest.SkipTest('cannot symlink')
+            try:
+                subprocess.run([
+                    sys.executable,
+                    '-mmypy',
+                    '--follow-imports', 'silent',
+                    '--check-untyped-defs',
+                    os.path.abspath(fn)],
+                    cwd=tmp_dir,
+                    check=True)
+            except subprocess.CalledProcessError as e:
+                raise AssertionError("mypy failed.  Look above this error for mypy's output.")
+
+
+if __name__ == '__main__':
+    run_tests()

tools/autograd/gen_python_functions.py

@@ -182,51 +182,75 @@ def should_generate_python_binding(declaration):
     return True
 
 
-def gen_py_variable_methods(out, declarations, template_path):
-    PY_VARIABLE_METHODS_CPP = CodeTemplate.from_file(template_path + '/python_variable_methods.cpp')
-    PY_VARIABLE_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_variable_methods_dispatch.h')
-
+def get_py_variable_methods(declarations):
+    """
+    Get declarations (grouped by name) which should be generated
+    as methods on Tensor.
+    """
     def should_bind(declaration):
         return (should_generate_python_binding(declaration) and
                 declaration['mode'] != 'NN' and
                 declaration.get('python_module') != 'nn' and
                 'Tensor' in declaration['method_of'])
 
-    py_variable_methods = group_declarations_by_name(declarations, should_bind)
+    return group_declarations_by_name(declarations, should_bind)
+
+
+def gen_py_variable_methods(out, declarations, template_path):
+    PY_VARIABLE_METHODS_CPP = CodeTemplate.from_file(template_path + '/python_variable_methods.cpp')
+    PY_VARIABLE_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_variable_methods_dispatch.h')
+
+    py_variable_methods = get_py_variable_methods(declarations)
 
     env = create_python_bindings(py_variable_methods, True)
     write(out, 'python_variable_methods.cpp', PY_VARIABLE_METHODS_CPP, env)
     write(out, 'python_variable_methods_dispatch.h', PY_VARIABLE_DISPATCH_H, env)
 
 
+def get_py_nn_functions(declarations):
+    """
+    Get declarations (grouped by name) which should be generated
+    as functions in the "nn" module.
+    """
+    def should_bind(declaration):
+        return (should_generate_python_binding(declaration) and
+                (declaration['mode'] == 'NN' or declaration.get('python_module') == 'nn'))
+
+    return group_declarations_by_name(declarations, should_bind)
+
+
 def gen_py_nn_functions(out, declarations, template_path):
     PY_NN_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_nn_functions.cpp')
     PY_NN_FUNCTIONS_H = CodeTemplate.from_file(template_path + '/python_nn_functions.h')
     PY_NN_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_nn_functions_dispatch.h')
 
-    def should_bind(declaration):
-        return (should_generate_python_binding(declaration) and
-                (declaration['mode'] == 'NN' or declaration.get('python_module') == 'nn'))
-
-    py_nn_functions = group_declarations_by_name(declarations, should_bind)
+    py_nn_functions = get_py_nn_functions(declarations)
 
     env = create_python_bindings(py_nn_functions, has_self=False, is_module=True)
     write(out, 'python_nn_functions.cpp', PY_NN_FUNCTIONS_CPP, env)
     write(out, 'python_nn_functions.h', PY_NN_FUNCTIONS_H, env)
     write(out, 'python_nn_functions_dispatch.h', PY_NN_DISPATCH_H, env)
 
 
-def gen_py_torch_functions(out, declarations, template_path):
-    PY_TORCH_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_torch_functions.cpp')
-    PY_TORCH_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_torch_functions_dispatch.h')
-
+def get_py_torch_functions(declarations):
+    """
+    Get declarations (grouped by name) which should be generated
+    as functions in the "torch" module.
+    """
     def should_bind(declaration):
         return (should_generate_python_binding(declaration) and
                 declaration['mode'] != 'NN' and
                 declaration.get('python_module') != 'nn' and
                 'namespace' in declaration['method_of'])
 
-    py_torch_functions = group_declarations_by_name(declarations, should_bind)
+    return group_declarations_by_name(declarations, should_bind)
+
+
+def gen_py_torch_functions(out, declarations, template_path):
+    PY_TORCH_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_torch_functions.cpp')
+    PY_TORCH_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_torch_functions_dispatch.h')
+
+    py_torch_functions = get_py_torch_functions(declarations)
 
     env = create_python_bindings(py_torch_functions, has_self=False)
     write(out, 'python_torch_functions.cpp', PY_TORCH_FUNCTIONS_CPP, env)
@@ -800,7 +824,16 @@ def is_smaller(d1, d2):
 
 
 def get_python_signature(declaration, include_out):
-    # Compute the Python function signature for argument parsing
+    # Compute the Python function signature for argument parsing,
+    # as specified in torch/csrc/utils/python_arg_parser.h.  WARNING:
+    # this is NOT the same type signature as specified by PEP 484
+    # as understood by mypy; our format was independently developed
+    # and has some quirks to make it more suitable specifically
+    # for error parsing.
+    #
+    # For a translation to mypy-valid type signatures, see
+    # tools/gen_pyi.py.  If you change any logic here, please
+    # check that file too.
     py_formal_args = []
     output_args = []
     type_args = []

tools/autograd/utils.py

@@ -14,6 +14,12 @@
     from tools.shared.module_loader import import_module
     CodeTemplate = import_module('code_template', 'aten/src/ATen/code_template.py').CodeTemplate
 
+# You should use these lines, rather than doing it manually.
+# Especially if you see this error!
+#
+#     File "/usr/local/lib/python2.7/dist-packages/yaml/__init__.py", line 69, in load
+#       loader = Loader(stream)
+#     TypeError: 'module' object is not callable
 try:
     # use faster C loader if available
     from yaml import CLoader as YamlLoader