Fix type hints for None constants (pytorch#23029)

Summary:
The type hint was being ignored when emitting `None` constants, this also de-dups some testing code
](https://our.intern.facebook.com/intern/diff/16364572/)
Pull Request resolved: pytorch#23029

Pulled By: driazati

Differential Revision: D16364572

fbshipit-source-id: 64f3abd3e37ee49c209480a85ed4f1b8802e5d93

Author: davidriazati

Diff for: test/jit_utils.py

@@ -19,11 +19,14 @@
 from contextlib import contextmanager
 from functools import reduce
 from itertools import chain
+from torch._six import StringIO
+
 import inspect
 import io
 import math
 import os
 import pickle
+import sys
 import tempfile
 import textwrap
 
@@ -39,6 +42,21 @@ class JitTestCase(TestCase):
     _do_cuda_memory_leak_check = True
     _restored_warnings = False
 
+    class capture_stdout(list):
+        """
+        Replace sys.stdout with a temporary StringIO
+        """
+        def __enter__(self):
+            self.sys_stdout = sys.stdout
+            self.stringio = StringIO()
+            sys.stdout = self.stringio
+            return self
+
+        def __exit__(self, *args):
+            self.append(str(self.stringio.getvalue()))
+            del self.stringio
+            sys.stdout = self.sys_stdout
+
     def setHooks(self):
         torch._C._jit_set_emit_hooks(self.emitModuleHook, self.emitFunctionHook)
 

Diff for: test/test_jit.py

@@ -2850,21 +2850,6 @@ def test_print_op_module(self):
 
 
 class TestScript(JitTestCase):
-    class capture_stdout(list):
-        """
-        Replace sys.stdout with a temporary StringIO
-        """
-        def __enter__(self):
-            self.sys_stdout = sys.stdout
-            self.stringio = StringIO()
-            sys.stdout = self.stringio
-            return self
-
-        def __exit__(self, *args):
-            self.append(str(self.stringio.getvalue()))
-            del self.stringio
-            sys.stdout = self.sys_stdout
-
     def test_sequence_parsing(self):
         tests = [
             ("return [x, x,]", True),
@@ -3194,35 +3179,6 @@ def annotate_none_no_optional():
         self.checkScript(annotate_none, ())
         self.checkScript(annotate_none_no_optional, ())
 
-    @unittest.skipIf(True, "Python 3 required")
-    def test_type_annotate_py3(self):
-        code = dedent("""
-        import torch
-        def fn():
-            a : List[int] = []
-            b : torch.Tensor = torch.ones(2, 2)
-            for _ in range(10):
-                a.append(4)
-            return a, b
-        """)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            script_path = os.path.join(tmp_dir, 'script.py')
-            with open(script_path, 'w') as f:
-                f.write(code)
-            fn = get_fn('test_type_annotate_py3', script_path)
-
-            self.checkScript(fn, ())
-
-        code = dedent("""
-        def wrong_type():
-            wrong : List[int] = [0.5]
-            return wrong
-        """)
-
-        with self.assertRaisesRegex(RuntimeError, "Lists must contain only a single type"):
-            cu = torch.jit.CompilationUnit(code)
-
     def test_robust_op_resolution(self):
         neg = torch.add  # misleading name to make sure we resolve by function
 
@@ -7862,7 +7818,7 @@ def foo(i):
         v = torch.rand(10, 3)
         self.checkScript(foo, (v,))
 
-        with self.assertRaisesRegex(RuntimeError, r"variable 'a' previously has type Tuple"):
+        with self.assertRaisesRegex(RuntimeError, r"Variable 'a' previously has type Tuple"):
             @torch.jit.script
             def mixtypes(x):
                 a = (x, x)
@@ -7890,7 +7846,7 @@ def diff_type_used():
                     c0 = 1.0
                 return c0
 
-        with self.assertRaisesRegex(RuntimeError, "variable 'c0' previously has type float"):
+        with self.assertRaisesRegex(RuntimeError, "Variable 'c0' previously has type float"):
             @torch.jit.script
             def diff_existing_type(x):
                 c0 = 1.0

Diff for: test/test_jit_py3.py

@@ -1,54 +1,12 @@
-import sys
-import torch
-from torch.testing import FileCheck
 from common_utils import run_tests
-from contextlib import contextmanager
 from jit_utils import JitTestCase
+from torch.testing import FileCheck
 from typing import NamedTuple, List
 
-WINDOWS = sys.platform == 'win32'
+import torch
 
-class TestScriptPy3(JitTestCase):
-    @contextmanager
-    def capture_stdout(self):
-        # No idea how to capture stdout from C++ on Windows
-        if WINDOWS:
-            yield ['']
-            return
-        import os
-        import fcntl
-        import errno
-        sys.stdout.flush()
-        stdout_fd = os.dup(1)
-        r, w = os.pipe()
-        try:
-            # Override stdout with r - dup is guaranteed to return the lowest free fd
-            os.close(1)
-            os.dup(w)
-
-            captured_stdout = ['']
-            yield captured_stdout
-            sys.stdout.flush()  # Make sure that Python hasn't buffered anything
-
-            # Do the ugly dance to read all the data that was written into the pipe
-            fcntl.fcntl(r, fcntl.F_SETFL, os.O_NONBLOCK)
-            total_stdout = ''
-            while True:
-                try:
-                    total_stdout += os.read(r, 1000).decode('ascii')
-                except OSError as e:
-                    if e.errno != errno.EAGAIN:
-                        raise
-                    break
-            captured_stdout[0] = total_stdout
-        finally:
-            # Revert the change, and clean up all fds
-            os.close(1)
-            os.dup(stdout_fd)
-            os.close(stdout_fd)
-            os.close(r)
-            os.close(w)
 
+class TestScriptPy3(JitTestCase):
     def test_joined_str(self):
         def func(x):
             hello, test = "Hello", "test"
@@ -212,5 +170,25 @@ def forward(self):
         for name in ['a', 'b', 'c']:
             self.assertEqual(getattr(out_loaded, name), getattr(out, name))
 
+    def test_type_annotate_py3(self):
+        def fn():
+            a : List[int] = []
+            b : torch.Tensor = torch.ones(2, 2)
+            c : Optional[torch.Tensor] = None
+            for _ in range(10):
+                a.append(4)
+                c = torch.ones(2, 2)
+            return a, b, c
+
+        self.checkScript(fn, ())
+
+        def wrong_type():
+            wrong : List[int] = [0.5]
+            return wrong
+
+        with self.assertRaisesRegex(RuntimeError, "Lists must contain only a single type"):
+            torch.jit.script(wrong_type)
+
+
 if __name__ == '__main__':
     run_tests()

Diff for: torch/csrc/jit/constants.cpp

@@ -90,9 +90,20 @@ c10::optional<Value*> tryInsertConstant(
     n->setScope(*scope);
   if (result_type) {
     auto inferred_type = n->output()->type();
-    // Retain more type information in case of tensor constant
-    if (!(inferred_type->isSubtypeOf(TensorType::get()) &&
-          result_type->isSubtypeOf(inferred_type))) {
+
+    if (inferred_type->isSubtypeOf(NoneType::get()) &&
+        !inferred_type->isSubtypeOf(result_type)) {
+      // None doesn't subtype Optional, but an Optional can be None, so handle
+      // that here
+      if (result_type->kind() == TypeKind::OptionalType) {
+        n->output()->setType(result_type);
+      } else {
+        // Implicitly wrap non-optionals
+        n->output()->setType(OptionalType::create(result_type));
+      }
+    } else if (!(inferred_type->isSubtypeOf(TensorType::get()) &&
+                 result_type->isSubtypeOf(inferred_type))) {
+      // Retain more type information in case of tensor constant
       n->output()->setType(result_type);
     }
   }

Diff for: torch/csrc/jit/script/compiler.cpp

@@ -307,20 +307,21 @@ struct Environment {
       }
       if (!as_simple_value->type()->isSubtypeOf(
               unshapedType(simple_parent->type()))) {
-        std::stringstream errMsg;
-        errMsg << "variable '" << name << "' previously has type "
+        auto error = ErrorReport(loc);
+        error << "Variable '" << name << "' previously has type "
                << simple_parent->type()->python_str()
                << " but is now being assigned to a value of type "
                << as_simple_value->type()->python_str();
+
         // Special-cased error msg if we're trying to assign to a tensor list.
         if (simple_parent->type()->kind() == TypeKind::ListType &&
             as_simple_value->type()->kind() == TypeKind::ListType) {
-          errMsg << "\n. (Note: empty lists are constructed as Tensor[]; "
+          error << "\n. (Note: empty lists are constructed as Tensor[]; "
                  << "if you want an empty list of a different type, "
                  << "use `torch.jit.annotate(List[T], [])`, "
                  << "where `T` is the type of elements in the list)";
         }
-        throw ErrorReport(loc) << errMsg.str();
+        throw error;
       }
     }
     if (as_simple_value) {
@@ -2455,7 +2456,7 @@ struct to_ir {
         return graph->insertConstant(false, nullptr, tree->range());
       } break;
       case TK_NONE: {
-        return graph->insertConstant(IValue(), nullptr, tree->range());
+        return graph->insertConstant(IValue(), type_hint, tree->range());
       } break;
       case TK_SUBSCRIPT: {
         return emitSubscript(Subscript(tree));
@@ -2542,7 +2543,6 @@ struct to_ir {
       } break;
       default:
         throw ErrorReport(tree) << "Cannot emit expr for: " << tree;
-        break;
     }
   }
 
@@ -2694,7 +2694,11 @@ struct to_ir {
         ++dim;
         continue;
       }
-      auto index = emitExpr(subscript_expr, OptionalType::ofTensor());
+      TypePtr type_hint = OptionalType::ofTensor();
+      if (subscript_expr.kind() == TK_NONE) {
+        type_hint = NoneType::get();
+      }
+      auto index = emitExpr(subscript_expr, type_hint);
       if (index->type() == IntType::get()) {
         // NB: note, select squeezes out a dimension,
         // so dim is **not** incremented