feat: update Pipe to allow custom init in subclass

Author: flavioschneider

pipd/pipe.py

@@ -2,7 +2,7 @@
 
 import re
 import traceback
-from typing import Any, Callable, Dict, Iterator, Optional, Type
+from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Type
 
 
 def is_iterable(obj):
@@ -52,46 +52,63 @@ def method(*args, **kwargs):
 
 
 class Pipe(metaclass=PipeMeta):
-    functions: Dict[str, Type[Function]] = {}
+    _functions: Dict[str, Type[Function]] = {}
 
     def __init__(
         self,
         *iterable,
-        function: Optional[Callable] = identity,
-        handler: Optional[Callable] = log_traceback_and_continue,
+        function: Callable = identity,
+        handler: Callable = log_traceback_and_continue,
     ):
         if len(iterable) == 1 and is_iterable(iterable[0]):
             iterable = iterable[0]
-        self.iterable = iterable
-        self.function = function
-        self.handler = handler
+        self._iterable: Iterable = iterable
+        self._function = function
+        self._handler = handler
 
     def __iter__(self):
-        for item in self.function(iter(self.iterable)):
+        for item in self._function(iter(self._iterable)):
             try:
                 yield item
             except Exception as e:
-                self.handler(e)
+                self._handler(e)
 
-    def close(self):
-        # Necessary to use `yield from` on a Pipe object
-        pass
-
-    def list(self):
-        return list(self)
+    def _new_pipe(
+        self,
+        iterable: Optional[Iterable] = None,
+        function: Optional[Callable] = None,
+        functions: Optional[Dict[str, Type[Function]]] = None,
+        handler: Optional[Callable] = None,
+    ):
+        # This is for composition intead of init so that init can be customized
+        pipe = self.__class__()
+        pipe._iterable = iterable or self._iterable
+        pipe._function = function or self._function
+        pipe._functions = functions or self._functions
+        pipe._handler = handler or self._handler
+        return pipe
 
     def __getattr__(self, name):
         def method(*args, **kwargs):
-            cls = self.__class__
-            function = self.functions[name](*args, **kwargs)
-            return cls(self.iterable, function=compose(self.function, function))
+            return self._new_pipe(
+                function=compose(self._function, self._functions[name](*args, **kwargs))
+            )
 
         return method
 
     def __call__(self, *iterable):
-        return self.__class__(*iterable, function=self.function)
+        if len(iterable) == 1 and is_iterable(iterable[0]):
+            iterable = iterable[0]
+        return self._new_pipe(iterable=iterable)
+
+    def close(self):
+        # Necessary to use `yield from` on a Pipe object
+        pass
+
+    def list(self):
+        return list(self)
 
     @classmethod
     def add_fn(cls, fn: Type[Function], name: Optional[str] = None):
         fn_name = camelcase_to_snakecase(fn.__name__) if name is None else name
-        cls.functions[fn_name] = fn
+        cls._functions[fn_name] = fn

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="pipd",
     packages=find_packages(exclude=[]),
-    version="0.1.7",
+    version="0.1.8",
     description="Utility functions for python data pipelines.",
     long_description_content_type="text/markdown",
     author="ElevenLabs",