From f902d87092a265708626376202ab430beb75fcd8 Mon Sep 17 00:00:00 2001
From: Nicholas Gates <gatesn@users.noreply.github.com>
Date: Mon, 2 Oct 2023 19:37:11 +0100
Subject: [PATCH] Add a PyIter type (#168)

---
 pydust/src/builtins.zig   | 26 +++++++++++++-----
 pydust/src/types.zig      |  1 +
 pydust/src/types/iter.zig | 55 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 7 deletions(-)
 create mode 100644 pydust/src/types/iter.zig

diff --git a/pydust/src/builtins.zig b/pydust/src/builtins.zig
index c42b0b27..e788ff28 100644
--- a/pydust/src/builtins.zig
+++ b/pydust/src/builtins.zig
@@ -85,13 +85,6 @@ pub fn is_none(object: anytype) bool {
     return ffi.Py_IsNone(obj.py) == 1;
 }
 
-/// Get the length of the given object. Equivalent to len(obj) in Python.
-pub fn len(object: anytype) !usize {
-    const length = ffi.PyObject_Length(py.object(object).py);
-    if (length < 0) return PyError.PyRaised;
-    return @intCast(length);
-}
-
 /// Import a module by fully-qualified name returning a PyObject.
 pub fn import(module_name: [:0]const u8) !py.PyObject {
     return (try py.PyModule.import(module_name)).obj;
@@ -114,6 +107,25 @@ pub fn isinstance(object: anytype, cls: anytype) !bool {
     return result == 1;
 }
 
+/// Return an iterator for the given object if it has one. Equivalent to iter(obj) in Python.
+pub fn iter(object: anytype) !py.PyIter {
+    const iterator = ffi.PyObject_GetIter(py.object(object).py) orelse return PyError.PyRaised;
+    return py.PyIter.unchecked(.{ .py = iterator });
+}
+
+/// Get the length of the given object. Equivalent to len(obj) in Python.
+pub fn len(object: anytype) !usize {
+    const length = ffi.PyObject_Length(py.object(object).py);
+    if (length < 0) return PyError.PyRaised;
+    return @intCast(length);
+}
+
+/// Return the next item of an iterator. Equivalent to next(obj) in Python.
+pub fn next(comptime T: type, iterator: anytype) !?T {
+    const pyiter = try py.PyIter.checked(iterator);
+    return try pyiter.next(T);
+}
+
 /// Return "false" if the object is considered to be truthy, and true otherwise.
 pub fn not_(object: anytype) !bool {
     const result = ffi.PyObject_Not(py.object(object).py);
diff --git a/pydust/src/types.zig b/pydust/src/types.zig
index c2a7859f..7e06a30f 100644
--- a/pydust/src/types.zig
+++ b/pydust/src/types.zig
@@ -16,6 +16,7 @@ pub usingnamespace @import("types/bytes.zig");
 pub usingnamespace @import("types/dict.zig");
 pub usingnamespace @import("types/error.zig");
 pub usingnamespace @import("types/float.zig");
+pub usingnamespace @import("types/iter.zig");
 pub usingnamespace @import("types/list.zig");
 pub usingnamespace @import("types/long.zig");
 pub usingnamespace @import("types/module.zig");
diff --git a/pydust/src/types/iter.zig b/pydust/src/types/iter.zig
new file mode 100644
index 00000000..15846520
--- /dev/null
+++ b/pydust/src/types/iter.zig
@@ -0,0 +1,55 @@
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//         http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+const std = @import("std");
+const py = @import("../pydust.zig");
+const PyObjectMixin = @import("./obj.zig").PyObjectMixin;
+const ffi = py.ffi;
+const PyError = @import("../errors.zig").PyError;
+
+/// Wrapper for Python PyIter.
+/// Constructed using py.iter(...)
+pub const PyIter = extern struct {
+    obj: py.PyObject,
+
+    pub usingnamespace PyObjectMixin("iterator", "PyIter", @This());
+
+    pub fn next(self: PyIter, comptime T: type) !?T {
+        if (ffi.PyIter_Next(self.obj.py)) |result| {
+            return try py.as(T, result);
+        }
+
+        // If no exception, then the item is missing.
+        if (ffi.PyErr_Occurred() == null) {
+            return null;
+        }
+
+        return PyError.PyRaised;
+    }
+
+    // TODO(ngates): implement PyIter_Send when required
+};
+
+test "PyIter" {
+    py.initialize();
+    defer py.finalize();
+
+    const tuple = try py.PyTuple.create(.{ 1, 2, 3 });
+    defer tuple.decref();
+
+    const iterator = try py.iter(tuple);
+    var previous: u64 = 0;
+    while (try iterator.next(u64)) |v| {
+        try std.testing.expect(v > previous);
+        previous = v;
+    }
+}