diff --git a/stdlib/public/Python/Python.swift b/stdlib/public/Python/Python.swift index 3a4305e22a809..25cd2f63d5c7d 100644 --- a/stdlib/public/Python/Python.swift +++ b/stdlib/public/Python/Python.swift @@ -1257,6 +1257,29 @@ extension PythonObject : MutableCollection { } } +extension PythonObject : Sequence { + public struct Iterator : IteratorProtocol { + fileprivate let pythonIterator: PythonObject + + public func next() -> PythonObject? { + guard let result = PyIter_Next(self.pythonIterator.borrowedPyObject) else { + try! throwPythonErrorIfPresent() + return nil + } + return PythonObject(consuming: result) + } + } + + public func makeIterator() -> Iterator { + guard let result = PyObject_GetIter(borrowedPyObject) else { + try! throwPythonErrorIfPresent() + // Unreachable. A Python `TypeError` must have been thrown. + preconditionFailure() + } + return Iterator(pythonIterator: PythonObject(consuming: result)) + } +} + //===----------------------------------------------------------------------===// // `ExpressibleByLiteral` conformances //===----------------------------------------------------------------------===// diff --git a/stdlib/public/Python/PythonLibrary+Symbols.swift b/stdlib/public/Python/PythonLibrary+Symbols.swift index 22e334e6ca76b..1c3d590ee0df9 100644 --- a/stdlib/public/Python/PythonLibrary+Symbols.swift +++ b/stdlib/public/Python/PythonLibrary+Symbols.swift @@ -123,6 +123,14 @@ let PyDict_Next: @convention(c) ( UnsafeMutablePointer) -> Int32 = PythonLibrary.loadSymbol(name: "PyDict_Next") +let PyIter_Next: @convention(c) ( + PyObjectPointer) -> PyObjectPointer? = + PythonLibrary.loadSymbol(name: "PyIter_Next") + +let PyObject_GetIter: @convention(c) ( + PyObjectPointer) -> PyObjectPointer? = + PythonLibrary.loadSymbol(name: "PyObject_GetIter") + let PyList_New: @convention(c) (Int) -> PyObjectPointer? = PythonLibrary.loadSymbol(name: "PyList_New") diff --git a/test/Python/python_runtime.swift b/test/Python/python_runtime.swift index 1bf91712608d2..4e6d7e7d1101f 100644 --- a/test/Python/python_runtime.swift +++ b/test/Python/python_runtime.swift @@ -77,6 +77,14 @@ PythonRuntimeTestSuite.testWithLeakChecking("PythonDict") { expectEqual("d", dict["b"]) } +PythonRuntimeTestSuite.testWithLeakChecking("Iterator") { + var sum = PythonObject(0) + for v in Python.iter([1, 2, 3]) { + sum += v + } + expectEqual(6, sum) +} + PythonRuntimeTestSuite.testWithLeakChecking("Range") { let slice = PythonObject(5..<10) expectEqual(Python.slice(5, 10), slice)