Merge pull request #21 from caspervdw/pipeline-properties

ENH Property-changing pipelines

Author: danielballan

slicerator.py

@@ -519,7 +519,57 @@ def wrapper(f):
         return wrapper(func)
 
 
-def _pipeline(func, retain_doc=False):
+def _pipeline(func_or_class, **kwargs):
+    try:
+        is_class = issubclass(func_or_class, Pipeline)
+    except TypeError:
+        is_class = False
+    if is_class:
+        return _pipeline_fromclass(func_or_class, **kwargs)
+    else:
+        return _pipeline_fromfunc(func_or_class, **kwargs)
+
+
+def _pipeline_fromclass(cls, retain_doc=False):
+    """Actual `pipeline` implementation
+
+    Parameters
+    ----------
+    func : class
+        Class for lazy evaluation
+    retain_doc : bool
+        If True, don't modify `func`'s doc string to say that it has been
+        made lazy
+
+    Returns
+    -------
+    Pipeline
+        Lazy function evaluation :py:class:`Pipeline` for `func`.
+    """
+    @wraps(cls)
+    def process(obj, *args, **kwargs):
+        if hasattr(obj, '_slicerator_flag') or isinstance(obj, Slicerator) \
+                or isinstance(obj, Pipeline):
+            return cls(obj, *args, **kwargs)
+        else:
+            # Fall back on normal behavior of func, interpreting input
+            # as a single image.
+            return cls([obj], *args, **kwargs)[0]
+
+    if not retain_doc:
+        if process.__doc__ is None:
+            process.__doc__ = ''
+        process.__doc__ = ("This function has been made lazy. When passed\n"
+                           "a Slicerator, it will return a \n"
+                           "Pipeline of the results. When passed \n"
+                           "any other objects, its behavior is "
+                           "unchanged.\n\n") + process.__doc__
+    process.__name__ = cls.__name__
+    return process
+
+
+
+def _pipeline_fromfunc(func, retain_doc=False):
     """Actual `pipeline` implementation
 
     Parameters
@@ -541,6 +591,7 @@ def process(obj, *args, **kwargs):
                 or isinstance(obj, Pipeline):
             def proc_func(x):
                 return func(x, *args, **kwargs)
+
             return Pipeline(obj, proc_func)
         else:
             # Fall back on normal behavior of func, interpreting input

tests.py

@@ -9,7 +9,8 @@
 from six import BytesIO
 import pickle
 from nose.tools import assert_true, assert_false, assert_equal, assert_raises
-from slicerator import Slicerator, pipeline, index_attr, propagate_attr
+from numpy.testing import assert_array_equal
+from slicerator import Slicerator, Pipeline, pipeline, index_attr, propagate_attr
 
 path, _ = os.path.split(os.path.abspath(__file__))
 path = os.path.join(path, 'data')
@@ -350,6 +351,29 @@ def test_composed_pipelines():
     assert_letters_equal(composed, 'zbCdefghij')
 
 
+def test_pipeline_class():
+    sli = Slicerator(np.empty((10, 32, 64)))
+
+    @pipeline
+    class crop(Pipeline):
+        def __init__(self, reader, bbox):
+            self.bbox = bbox
+            Pipeline.__init__(self, reader, None)
+
+        def _get(self, key):
+            bbox = self.bbox
+            return self._ancestor[key][bbox[0]:bbox[2], bbox[1]:bbox[3]]
+
+        @property
+        def frame_shape(self):
+            bbox = self.bbox
+            return (bbox[2] - bbox[0], bbox[3] - bbox[1])
+
+    cropped = crop(sli, (5, 5, 10, 20))
+    assert_array_equal(cropped[0], sli[0][5:10, 5:20])
+    assert_array_equal(cropped.frame_shape, (5, 15))
+
+
 def test_serialize():
     # dump Slicerator
     stream = BytesIO()