Improved python io handling

mcneel · Aug 27, 2024 · a1724fb · a1724fb
1 parent 5790346
commit a1724fb
Showing 1 changed file with 99 additions and 35 deletions.
diff --git a/src/runtime/McNeel.PythonEngine.cs b/src/runtime/McNeel.PythonEngine.cs
@@ -266,22 +266,17 @@ public void SetSysArgs(IEnumerable<string> args)
         #endregion
 
         #region Standard IO
+        PyObject _defaultStdOut = default;
+        PyObject _defaultStdErr = default;
+
         public void SetStdOut(Stream stdout)
         {
             using (Py.GIL())
             {
                 using var sys = Runtime.PyImport_ImportModule("sys");
                 using (PyObject sysObj = sys.MoveToPyObject())
                 {
-                    if (sysObj.GetAttr("stdout") is PyObject stdoutObj
-                            && ManagedType.GetManagedObject(stdoutObj) is CLRObject stdoutClrObj
-                            && stdoutClrObj.inst is PythonStream exstStdout)
-                    {
-                        exstStdout.Dispose();
-                    }
-
-                    using var stdio = PyObject.FromManagedObject(new PythonStream(stdout, 1));
-                    sysObj.SetAttr("stdout", stdio);
+                    SetStdOutObj(sysObj, stdout);
                 }
             }
         }
@@ -293,15 +288,7 @@ public void SetStdErr(Stream stderr)
                 using var sys = Runtime.PyImport_ImportModule("sys");
                 using (PyObject sysObj = sys.MoveToPyObject())
                 {
-                    if (sysObj.GetAttr("stderr") is PyObject stderrObj
-                            && ManagedType.GetManagedObject(stderrObj) is CLRObject stderrClrObj
-                            && stderrClrObj.inst is PythonStream exstStdErr)
-                    {
-                        exstStdErr.Dispose();
-                    }
-
-                    using var newStderr = PyObject.FromManagedObject(new PythonStream(stderr, 2));
-                    sysObj.SetAttr("stderr", newStderr);
+                    SetStdErrObj(sysObj, stderr);
                 }
             }
         }
@@ -313,29 +300,106 @@ public void SetStdOutErr(Stream stdout, Stream stderr)
                 using var sys = Runtime.PyImport_ImportModule("sys");
                 using (PyObject sysObj = sys.MoveToPyObject())
                 {
-                    if (sysObj.GetAttr("stdout") is PyObject stdoutObj
-                            && ManagedType.GetManagedObject(stdoutObj) is CLRObject stdoutClrObj
-                            && stdoutClrObj.inst is PythonStream exstStdout)
-                    {
-                        exstStdout.Dispose();
-                    }
+                    SetStdOutObj(sysObj, stdout);
+                    SetStdErrObj(sysObj, stderr);
+                }
+            }
+        }
 
-                    if (sysObj.GetAttr("stderr") is PyObject stderrObj
-                            && ManagedType.GetManagedObject(stderrObj) is CLRObject stderrClrObj
-                            && stderrClrObj.inst is PythonStream exstStdErr)
-                    {
-                        exstStdErr.Dispose();
-                    }
+        public void ClearStdOut()
+        {
+            using (Py.GIL())
+            {
+                using var sys = Runtime.PyImport_ImportModule("sys");
+                using (PyObject sysObj = sys.MoveToPyObject())
+                {
+                    SetStdOutObj(sysObj);
+                }
+            }
+        }
+
+        public void ClearStdErr()
+        {
+            using (Py.GIL())
+            {
+                using var sys = Runtime.PyImport_ImportModule("sys");
+                using (PyObject sysObj = sys.MoveToPyObject())
+                {
+                    SetStdErrObj(sysObj);
+                }
+            }
+        }
+
+        public void ClearStdOutErr()
+        {
+            using (Py.GIL())
+            {
+                using var sys = Runtime.PyImport_ImportModule("sys");
+                using (PyObject sysObj = sys.MoveToPyObject())
+                {
+                    SetStdOutObj(sysObj);
+                    SetStdErrObj(sysObj);
+                }
+            }
+        }
+
+        void SetStdOutObj(PyObject sysObj, Stream stream = default)
+        {
+            if (sysObj.GetAttr("stdout") is PyObject stdoutObj)
+            {
+                if (ManagedType.GetManagedObject(stdoutObj) is CLRObject stdoutClrObj
+                        && stdoutClrObj.inst is PythonStream exstStdout)
+                {
+                    exstStdout.Dispose();
+                }
+                else if (_defaultStdOut is null)
+                {
+                    _defaultStdOut = stdoutObj;
+                }
+            }
 
-                    using var newStdout = PyObject.FromManagedObject(new PythonStream(stdout, 1));
-                    sysObj.SetAttr("stdout", newStdout);
-                    using var newStderr = PyObject.FromManagedObject(new PythonStream(stderr, 2));
-                    sysObj.SetAttr("stderr", newStderr);
+            if (stream is null)
+            {
+                if (_defaultStdOut is PyObject)
+                {
+                    sysObj.SetAttr("stdout", _defaultStdOut);
                 }
             }
+            else
+            {
+                using var stdio = PyObject.FromManagedObject(new PythonStream(stream, 1));
+                sysObj.SetAttr("stdout", stdio);
+            }
         }
 
-        public void ClearStdOutErr() => SetStdOutErr(default, default);
+        void SetStdErrObj(PyObject sysObj, Stream stream = default)
+        {
+            if (sysObj.GetAttr("stderr") is PyObject stderrObj)
+            {
+                if (ManagedType.GetManagedObject(stderrObj) is CLRObject stderrClrObj
+                        && stderrClrObj.inst is PythonStream exstStdErr)
+                {
+                    exstStdErr.Dispose();
+                }
+                else if (_defaultStdErr is null)
+                {
+                    _defaultStdErr = stderrObj;
+                }
+            }
+
+            if (stream is null)
+            {
+                if (_defaultStdErr is PyObject)
+                {
+                    sysObj.SetAttr("stderr", _defaultStdErr);
+                }
+            }
+            else
+            {
+                using var newStderr = PyObject.FromManagedObject(new PythonStream(stream, 2));
+                sysObj.SetAttr("stderr", newStderr);
+            }
+        }
         #endregion
 
         #region Execution