print a warning if a type annotation prefix is invalid according to mypy (#20884)

Krovatkin · facebook-github-bot · commit cbf2a4f5c413 · 2019-05-29T11:56:55.000-07:00
Summary: This PR adds a check that prints a warning if a type annotation prefix isn't what mypy expects. Pull Request resolved: #20884 Differential Revision: D15511043 Pulled By: Krovatkin fbshipit-source-id: 9038e074807832931faaa5f4e69628f94f51fd72
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -3373,6 +3373,28 @@ def __init__(self):
         self.assertEqual(D()(v), v + v)
 
 
+    def test_invalid_prefix_annotation(self):
+        with self.assertRaisesRegex(RuntimeError, "annotation prefix in line"):
+            with self.capture_stdout() as captured:
+                @torch.jit.script
+                def invalid_prefix_annotation1(a):
+                    #type: (Int) -> Int # noqa
+                    return a + 2
+
+        with self.assertRaisesRegex(RuntimeError, "annotation prefix in line"):
+            with self.capture_stdout() as captured:
+                @torch.jit.script
+                def invalid_prefix_annotation2(a):
+                    #type   : (Int) -> Int # noqa
+                    return a + 2
+
+        with self.assertRaisesRegex(RuntimeError, "annotation prefix in line"):
+            with self.capture_stdout() as captured:
+                @torch.jit.script
+                def invalid_prefix_annotation3(a):
+                    #     type: (Int) -> Int
+                    return a + 2
+
     def test_tracing_multiple_methods(self):
         class Net(nn.Module):
             def __init__(self):
diff --git a/torch/jit/annotations.py b/torch/jit/annotations.py
@@ -1,6 +1,7 @@
 import sys
 import ast
 import inspect
+import re
 import torch
 from .._jit_internal import List, BroadcastingList1, BroadcastingList2, \
     BroadcastingList3, Tuple, is_tuple, is_list, Dict, is_dict, Optional, \
@@ -117,8 +118,17 @@ def get_type_line(source):
     lines = source.split('\n')
     lines = [(line_num, line) for line_num, line in enumerate(lines)]
     type_lines = list(filter(lambda line: type_comment in line[1], lines))
+    lines_with_type = list(filter(lambda line: 'type' in line[1], lines))
+
 
     if len(type_lines) == 0:
+        type_pattern = re.compile('#[\t ]*type[\t ]*:')
+        wrong_type_lines = list(filter(lambda line: type_pattern.search(line[1]), lines))
+        if len(wrong_type_lines) > 0:
+            raise RuntimeError("The annotation prefix in line " + str(wrong_type_lines[0][0])
+                               + " is probably invalid.\nIt must be '# type:'"
+                               + "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)" # noqa
+                               + "\nfor examples")
         return None
     elif len(type_lines) == 1:
         # Only 1 type line, quit now
