fuck

James Reed · James Reed · commit 90c0674e00cd · 2019-06-11T14:32:08.000-07:00
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12934,16 +12934,22 @@ def test_sys_stdout_override(self):
         def foo():
             print('foo')
 
+        class Redirect(object):
+            def __init__(self):
+                self.s = ''
+
+            def write(self, s : str):
+                self.s += s
+
         old_stdout = sys.stdout
-        with tempfile.TemporaryFile() as f:
-            try:
-                sys.stdout = f
-                foo()
-            finally:
-                sys.stdout = old_stdout
+        redirect = Redirect()
+        try:
+            sys.stdout = redirect
+            foo()
+        finally:
+            sys.stdout = old_stdout
 
-            f.seek(0)
-            FileCheck().check('foo').run(f.read())
+        FileCheck().check('foo').run(redirect.s)
 
     def test_optional_tuple(self):
         def fn(x=None):
diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp
@@ -569,8 +569,17 @@ 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")(py::bytes(str));
+    try {
+      auto _stdout = py::module::import("sys").attr("stdout");
+      _stdout.attr("write")(str);
+    } catch (pybind11::error_already_set e) {
+      TORCH_WARN(
+          "Exception occured when printing to Python stdout: ",
+          e.what(),
+          ". Falling back to C++ stdout."
+          " Please report a bug to PyTorch");
+      std::cout << str << std::endl;
+    }
   });
 }
 } // namespace jit
diff --git a/torch/csrc/jit/print_handler.cpp b/torch/csrc/jit/print_handler.cpp
@@ -1,6 +1,7 @@
 #include <torch/csrc/jit/print_handler.h>
 
 #include <iostream>
+#include <string>
 
 namespace torch {
 namespace jit {
