switch

James Reed · James Reed · commit d089ea3417e4 · 2019-06-11T13:28:06.000-07:00
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12929,6 +12929,22 @@ def forward(self, key):
                 .check_count(s2, 1, exactly=True) \
                 .check_count("BINGET", 2, exactly=True).run(out.getvalue())
 
+    def test_sys_stdout_override(self):
+        @torch.jit.script
+        def foo():
+            print('foo')
+
+        old_stdout = sys.stdout
+        with tempfile.TemporaryFile() as f:
+            try:
+                sys.stdout = f
+                foo()
+            finally:
+                sys.stdout = old_stdout
+
+            f.seek(0)
+            FileCheck().check('foo').run(f.read())
+
     def test_optional_tuple(self):
         def fn(x=None):
             # type: (Optional[Tuple[int, int]]) -> Tuple[int, int]
diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp
@@ -54,8 +54,8 @@
 
 #include <ATen/core/function_schema.h>
 
-#include <Python.h>
 #include <pybind11/functional.h>
+#include <pybind11/iostream.h>
 
 #include <memory>
 #include <sstream>
@@ -570,7 +570,7 @@ void initJITBindings(PyObject* module) {
   setPrintHandler([](const std::string& str) {
     py::gil_scoped_acquire acquire;
     auto _stdout = py::module::import("sys").attr("stdout");
-    _stdout.attr("write")(str);
+    _stdout.attr("write")(py::bytes(str));
   });
 }
 } // namespace jit
diff --git a/torch/csrc/jit/init.h b/torch/csrc/jit/init.h
@@ -1,8 +1,5 @@
 #pragma once
 
-#include <functional>
-#include <iostream>
-
 namespace torch {
 namespace jit {
 
diff --git a/torch/csrc/jit/print_handler.cpp b/torch/csrc/jit/print_handler.cpp
@@ -1,5 +1,7 @@
 #include <torch/csrc/jit/print_handler.h>
 
+#include <iostream>
+
 namespace torch {
 namespace jit {
 
