API change: replace Pipeline.__iter__ with iterator(chunksize, prefetch) and add new PipelineIterator class

Author: gsakkis

pdal/libpdalpython.pyx

@@ -1,8 +1,10 @@
 # distutils: language = c++
 # cython: c_string_type=unicode, c_string_encoding=utf8
 
+import json
 from types import SimpleNamespace
 
+cimport cython
 from cython.operator cimport dereference as deref
 from cpython.ref cimport Py_DECREF
 from libcpp cimport bool
@@ -121,12 +123,33 @@ cdef extern from "PyPipeline.hpp" namespace "pdal::python":
     np.PyArrayObject* meshToNumpyArray(const TriangularMesh*) except +
 
 
-cdef class Pipeline:
-    cdef unique_ptr[StreamableExecutor] _executor
+@cython.internal
+cdef class PipelineResultsMixin:
+    @property
+    def log(self):
+        return self._get_executor().getLog()
+
+    @property
+    def schema(self):
+        return json.loads(self._get_executor().getSchema())
+
+    @property
+    def pipeline(self):
+        return self._get_executor().getPipeline()
+
+    @property
+    def metadata(self):
+        return self._get_executor().getMetadata()
+
+    cdef PipelineExecutor* _get_executor(self) except NULL:
+        raise NotImplementedError("Abstract method")
+
+
+
+cdef class Pipeline(PipelineResultsMixin):
+    cdef unique_ptr[PipelineExecutor] _executor
     cdef vector[shared_ptr[Array]] _inputs
     cdef int _loglevel
-    cdef int _chunk_size
-    cdef int _prefetch
 
     #========= writeable properties to be set before execution ===========================
 
@@ -150,44 +173,8 @@ cdef class Pipeline:
         self._loglevel = value
         self._del_executor()
 
-    @property
-    def chunk_size(self):
-        return self._chunk_size
-
-    @chunk_size.setter
-    def chunk_size(self, value):
-        assert value > 0
-        self._chunk_size = value
-        self._del_executor()
-
-    @property
-    def prefetch(self):
-        return self._prefetch
-
-    @prefetch.setter
-    def prefetch(self, value):
-        assert value >= 0
-        self._prefetch = value
-        self._del_executor()
-
     #========= readable properties to be read after execution ============================
 
-    @property
-    def log(self):
-        return self._get_executor().getLog()
-
-    @property
-    def schema(self):
-        return self._get_executor().getSchema()
-
-    @property
-    def pipeline(self):
-        return self._get_executor().getPipeline()
-
-    @property
-    def metadata(self):
-        return self._get_executor().getMetadata()
-
     @property
     def arrays(self):
         cdef PipelineExecutor* executor = self._get_executor()
@@ -215,18 +202,14 @@ cdef class Pipeline:
     def execute(self):
         return self._get_executor().execute()
 
-    def __iter__(self):
-        cdef StreamableExecutor* executor = self._get_executor()
-        try:
-            while True:
-                arr_ptr = executor.executeNext()
-                if arr_ptr is NULL:
-                    break
-                arr = <object>arr_ptr
-                Py_DECREF(arr)
-                yield arr
-        finally:
-            executor.stop()
+    def iterator(self, int chunk_size=10000, int prefetch = 0):
+        cdef StreamableExecutor* executor = new StreamableExecutor(
+            self._get_json(), chunk_size, prefetch
+        )
+        self._configure_executor(executor)
+        it = PipelineIterator()
+        it.set_executor(executor)
+        return it
 
     #========= non-public properties & methods ===========================================
 
@@ -243,12 +226,44 @@ cdef class Pipeline:
     def _del_executor(self):
         self._executor.reset()
 
-    cdef StreamableExecutor* _get_executor(self) except NULL:
+    cdef PipelineExecutor* _get_executor(self) except NULL:
         if not self._executor:
-            executor = new StreamableExecutor(self._get_json(),
-                                              self._chunk_size, self._prefetch)
-            executor.setLogLevel(self._loglevel)
-            executor.read()
-            addArrayReaders(executor, self._inputs)
+            executor = new PipelineExecutor(self._get_json())
+            self._configure_executor(executor)
             self._executor.reset(executor)
         return self._executor.get()
+
+    cdef _configure_executor(self, PipelineExecutor* executor):
+        executor.setLogLevel(self._loglevel)
+        executor.read()
+        addArrayReaders(executor, self._inputs)
+
+
+cdef class PipelineIterator(PipelineResultsMixin):
+    cdef unique_ptr[StreamableExecutor] _executor
+
+    cdef set_executor(self, StreamableExecutor* executor):
+        self._executor.reset(executor)
+
+    def __dealloc__(self):
+        self._executor.get().stop()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        cdef StreamableExecutor* executor = self._executor.get()
+        if executor.executed():
+            raise StopIteration
+
+        arr_ptr = executor.executeNext()
+        if arr_ptr is NULL:
+            executor.stop()
+            raise StopIteration
+
+        arr = <object> arr_ptr
+        Py_DECREF(arr)
+        return arr
+
+    cdef PipelineExecutor* _get_executor(self) except NULL:
+        return self._executor.get()

pdal/pipeline.py

@@ -28,8 +28,6 @@ def __init__(
         spec: Union[None, str, Sequence[Stage]] = None,
         arrays: Sequence[np.ndarray] = (),
         loglevel: int = logging.ERROR,
-        chunk_size: int = 10000,
-        prefetch: int = 0,
     ):
         super().__init__()
         self._stages: List[Stage] = []
@@ -39,8 +37,6 @@ def __init__(
                 self |= stage
         self.inputs = arrays
         self.loglevel = loglevel
-        self.chunk_size = chunk_size
-        self.prefetch = prefetch
 
     @property
     def stages(self) -> List[Stage]:
@@ -50,10 +46,6 @@ def stages(self) -> List[Stage]:
     def streamable(self) -> bool:
         return all(stage.streamable for stage in self._stages)
 
-    @property
-    def schema(self) -> Any:
-        return json.loads(super().schema)
-
     @property
     def loglevel(self) -> int:
         return LogLevelFromPDAL[super().loglevel]

test/test_pipeline.py

@@ -2,7 +2,6 @@
 import logging
 import os
 import sys
-from functools import partial
 
 import numpy as np
 import pytest
@@ -12,16 +11,12 @@
 DATADIRECTORY = os.path.join(os.path.dirname(__file__), "data")
 
 
-def get_pipeline(filename, *, chunk_size=None, prefetch=None):
+def get_pipeline(filename):
     with open(os.path.join(DATADIRECTORY, filename), "r") as f:
         if filename.endswith(".json"):
             pipeline = pdal.Pipeline(f.read())
         elif filename.endswith(".py"):
             pipeline = eval(f.read(), vars(pdal))
-    if chunk_size is not None:
-        pipeline.chunk_size = chunk_size
-    if prefetch is not None:
-        pipeline.prefetch = prefetch
     return pipeline
 
 
@@ -36,11 +31,6 @@ class TestPipeline:
     def test_construction(self, filename):
         """Can we construct a PDAL pipeline"""
         assert isinstance(get_pipeline(filename), pdal.Pipeline)
-        assert isinstance(get_pipeline(filename, chunk_size=100), pdal.Pipeline)
-        assert isinstance(get_pipeline(filename, prefetch=3), pdal.Pipeline)
-        assert isinstance(
-            get_pipeline(filename, chunk_size=100, prefetch=3), pdal.Pipeline
-        )
 
         # construct Pipeline from a sequence of stages
         r = pdal.Reader("r")
@@ -421,48 +411,51 @@ class TestPipelineIterator:
 
     def test_array(self):
         """Can we fetch PDAL data as numpy arrays"""
-        ri = get_pipeline("range.json", chunk_size=100)
-        arrays = list(ri)
-        assert len(arrays) == 11
-        concat_array = np.concatenate(arrays)
-
         r = get_pipeline("range.json")
         count = r.execute()
         arrays = r.arrays
         assert len(arrays) == 1
         array = arrays[0]
         assert count == len(array)
 
-        np.testing.assert_array_equal(array, concat_array)
+        for _ in range(10):
+            arrays = list(r.iterator(chunk_size=100))
+            assert len(arrays) == 11
+            concat_array = np.concatenate(arrays)
+            np.testing.assert_array_equal(array, concat_array)
+
+    def test_StopIteration(self):
+        """Is StopIteration raised when the iterator is exhausted"""
+        r = get_pipeline("range.json")
+        it = r.iterator(chunk_size=100)
+        for array in it:
+            assert isinstance(array, np.ndarray)
+        with pytest.raises(StopIteration):
+            next(it)
+        assert next(it, None) is None
 
     def test_metadata(self):
         """Can we fetch PDAL metadata"""
-        ri = get_pipeline("range.json", chunk_size=100)
-        with pytest.raises(RuntimeError):
-            ri.metadata
-        list(ri)
-
         r = get_pipeline("range.json")
-        with pytest.raises(RuntimeError):
-            r.metadata
         r.execute()
 
-        assert ri.metadata == r.metadata
+        it = r.iterator(chunk_size=100)
+        for _ in it:
+            pass
+
+        assert r.metadata == it.metadata
 
     @pytest.mark.xfail
     def test_schema(self):
         """Fetching a schema works"""
-        ri = get_pipeline("range.json", chunk_size=100)
-        with pytest.raises(RuntimeError):
-            ri.schema
-        list(ri)
-
         r = get_pipeline("range.json")
-        with pytest.raises(RuntimeError):
-            r.schema
         r.execute()
 
-        assert ri.schema == r.schema
+        it = r.iterator(chunk_size=100)
+        for _ in it:
+            pass
+
+        assert r.schema == it.schema
 
     def test_merged_arrays(self):
         """Can we load data from a list of arrays to PDAL"""
@@ -480,8 +473,7 @@ def test_merged_arrays(self):
         p.execute()
         non_streaming_array = np.concatenate(p.arrays)
         for chunk_size in range(5, 100, 5):
-            streaming_arrays = list(pdal.Pipeline(filter_intensity, arrays,
-                                                  chunk_size=chunk_size))
+            streaming_arrays = list(p.iterator(chunk_size=chunk_size))
             np.testing.assert_array_equal(np.concatenate(streaming_arrays),
                                           non_streaming_array)
 
@@ -492,7 +484,17 @@ def test_premature_exit(self):
         assert len(r.arrays) == 1
         array = r.arrays[0]
 
-        ri = get_pipeline("range.json", chunk_size=100)
-        for array2 in ri:
-            np.testing.assert_array_equal(array2, array[:len(array2)])
-            break
+        for _ in range(10):
+            for array2 in r.iterator(chunk_size=100):
+                np.testing.assert_array_equal(array2, array[:len(array2)])
+                break
+
+    def test_multiple_iterators(self):
+        """Can we create multiple independent iterators"""
+        r = get_pipeline("range.json")
+        it1 = r.iterator(chunk_size=100)
+        it2 = r.iterator(chunk_size=100)
+        for a1, a2 in zip(it1, it2):
+            np.testing.assert_array_equal(a1, a2)
+        assert next(it1, None) is None
+        assert next(it2, None) is None