From cebc51536201bcb4173046d93b26e4f5f3047fd2 Mon Sep 17 00:00:00 2001
From: Victor Milovanov <lost@losttech.software>
Date: Thu, 17 Dec 2020 14:47:29 -0800
Subject: [PATCH] PyIter: do not force dispose previous object upon moving to
 the next one

also added exception handling for PyIter_Next

PyObject now implements `IEnumerable<PyObject>`
---
 CHANGELOG.md                  |  2 ++
 src/embed_tests/TestPyIter.cs | 37 +++++++++++++++++++++++++++++++++++
 src/runtime/pyiter.cs         | 34 +++++++++++++-------------------
 src/runtime/pyobject.cs       |  5 +++--
 src/runtime/runtime.cs        |  2 ++
 5 files changed, 58 insertions(+), 22 deletions(-)
 create mode 100644 src/embed_tests/TestPyIter.cs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9afccbf2b..27df6dcbe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,7 @@ details about the cause of the failure
 -   BREAKING: Parameters marked with `ParameterAttributes.Out` are no longer returned in addition
      to the regular method return value (unless they are passed with `ref` or `out` keyword).
 -   BREAKING: Drop support for the long-deprecated CLR.* prefix.
+-   `PyObject` now implements `IEnumerable<PyObject>` in addition to `IEnumerable`
 
 ### Fixed
 
@@ -40,6 +41,7 @@ details about the cause of the failure
 -    Fixed a bug where indexers could not be used if they were inherited
 -    Made it possible to use `__len__` also on `ICollection<>` interface objects
 -    Made it possible to call `ToString`, `GetHashCode`, and `GetType` on inteface objects
+-    Fixed objects returned by enumerating `PyObject` being disposed too soon
 
 ### Removed
 
diff --git a/src/embed_tests/TestPyIter.cs b/src/embed_tests/TestPyIter.cs
new file mode 100644
index 000000000..7428da979
--- /dev/null
+++ b/src/embed_tests/TestPyIter.cs
@@ -0,0 +1,37 @@
+using System.Linq;
+using System.Text;
+
+using NUnit.Framework;
+
+using Python.Runtime;
+
+namespace Python.EmbeddingTest
+{
+    class TestPyIter
+    {
+        [OneTimeSetUp]
+        public void SetUp()
+        {
+            PythonEngine.Initialize();
+        }
+
+        [OneTimeTearDown]
+        public void Dispose()
+        {
+            PythonEngine.Shutdown();
+        }
+
+        [Test]
+        public void KeepOldObjects()
+        {
+            using (Py.GIL())
+            using (var testString = new PyString("hello world! !$%&/()=?"))
+            {
+                PyObject[] chars = testString.ToArray();
+                Assert.IsTrue(chars.Length > 1);
+                string reconstructed = string.Concat(chars.Select(c => c.As<string>()));
+                Assert.AreEqual(testString.As<string>(), reconstructed);
+            }
+        }
+    }
+}
diff --git a/src/runtime/pyiter.cs b/src/runtime/pyiter.cs
index 4ad761db7..2016ef4f8 100644
--- a/src/runtime/pyiter.cs
+++ b/src/runtime/pyiter.cs
@@ -9,7 +9,7 @@ namespace Python.Runtime
     /// PY3: https://docs.python.org/3/c-api/iterator.html
     /// for details.
     /// </summary>
-    public class PyIter : PyObject, IEnumerator<object>
+    public class PyIter : PyObject, IEnumerator<PyObject>
     {
         private PyObject _current;
 
@@ -46,41 +46,35 @@ public static PyIter GetIter(PyObject iterable)
 
         protected override void Dispose(bool disposing)
         {
-            if (null != _current)
-            {
-                _current.Dispose();
-                _current = null;
-            }
+            _current = null;
             base.Dispose(disposing);
         }
 
         public bool MoveNext()
         {
-            // dispose of the previous object, if there was one
-            if (null != _current)
+            NewReference next = Runtime.PyIter_Next(Reference);
+            if (next.IsNull())
             {
-                _current.Dispose();
-                _current = null;
-            }
+                if (Exceptions.ErrorOccurred())
+                {
+                    throw new PythonException();
+                }
 
-            IntPtr next = Runtime.PyIter_Next(obj);
-            if (next == IntPtr.Zero)
-            {
+                // stop holding the previous object, if there was one
+                _current = null;
                 return false;
             }
 
-            _current = new PyObject(next);
+            _current = next.MoveToPyObject();
             return true;
         }
 
         public void Reset()
         {
-            //Not supported in python.
+            throw new NotSupportedException();
         }
 
-        public object Current
-        {
-            get { return _current; }
-        }
+        public PyObject Current => _current;
+        object System.Collections.IEnumerator.Current => _current;
     }
 }
diff --git a/src/runtime/pyobject.cs b/src/runtime/pyobject.cs
index 6957db8df..524d6e4b6 100644
--- a/src/runtime/pyobject.cs
+++ b/src/runtime/pyobject.cs
@@ -16,7 +16,7 @@ namespace Python.Runtime
     /// for details.
     /// </summary>
     [Serializable]
-    public partial class PyObject : DynamicObject, IEnumerable, IDisposable
+    public partial class PyObject : DynamicObject, IEnumerable<PyObject>, IDisposable
     {
 #if TRACE_ALLOC
         /// <summary>
@@ -705,10 +705,11 @@ public PyObject GetIterator()
         /// python object to be iterated over in C#. A PythonException will be
         /// raised if the object is not iterable.
         /// </remarks>
-        public IEnumerator GetEnumerator()
+        public IEnumerator<PyObject> GetEnumerator()
         {
             return PyIter.GetIter(this);
         }
+        IEnumerator IEnumerable.GetEnumerator() => this.GetEnumerator();
 
 
         /// <summary>
diff --git a/src/runtime/runtime.cs b/src/runtime/runtime.cs
index 10cbdc869..9204f04b6 100644
--- a/src/runtime/runtime.cs
+++ b/src/runtime/runtime.cs
@@ -1854,6 +1854,8 @@ internal static bool PyIter_Check(IntPtr pointer)
 
         [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
         internal static extern IntPtr PyIter_Next(IntPtr pointer);
+        [DllImport(_PythonDll, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern NewReference PyIter_Next(BorrowedReference pointer);
 
 
         //====================================================================