소스 검색

startup: add boolean kwarg "catch" to PipeMap & PipePair

Fabian Peter Hammerle 3 일 전
부모
커밋
812bb29555
1개의 변경된 파일34개의 추가작업 그리고 4개의 파일을 삭제
  1. 34 4
      profile_default/startup/init.py

+ 34 - 4
profile_default/startup/init.py

@@ -178,10 +178,25 @@ class PipeMap(Pipe):
             return functools.partial(map, function)
         return functools.partial(map, cls._partial_map(function, axis=axis - 1))
 
+    @staticmethod
+    def _catch(
+        *, function: typing.Callable[[typing.Any], typing.Any], arg: typing.Any
+    ) -> typing.Any | Exception:
+        try:
+            return function(arg)
+        except Exception as exc:
+            return exc
+
     def __init__(
-        self, function: typing.Callable[[typing.Any], typing.Any], axis: int = 0
+        self,
+        function: typing.Callable[[typing.Any], typing.Any],
+        axis: int = 0,
+        catch: bool = False,
     ) -> None:
-        self._function = self._partial_map(function, axis=axis)
+        self._function = self._partial_map(
+            (lambda v: self._catch(function=function, arg=v)) if catch else function,
+            axis=axis,
+        )
 
 
 assert list(PipeMap._partial_map(str, axis=0)(range(3))) == ["0", "1", "2"]
@@ -202,6 +217,9 @@ assert "123|456\n98|76|54".splitlines() | PipeMap(lambda s: s.split("|")) | Pipe
     ((1, 2, 3), (4, 5, 6)),
     ((9, 8), (7, 6), (5, 4)),
 ]
+assert [1, 0, 2] | PipeMap(lambda d: 42 // d, catch=True) | PipeMap(
+    lambda r: r.args if isinstance(r, Exception) else r
+) | Pipe(list) == [42, ("integer division or modulo by zero",), 21]
 
 
 class PipeFilter(Pipe):
@@ -218,10 +236,19 @@ assert [0, True, 2.3, "4", (5, 6)] | PipeFilter(
 
 
 class PipePair(PipeMap):
+    @staticmethod
+    def _catch(
+        *, function: typing.Callable[[typing.Any], typing.Any], arg: typing.Any
+    ) -> typing.Tuple[typing.Any, typing.Any | Exception]:
+        try:
+            return function(arg)
+        except Exception as exc:
+            return (arg, exc)
+
     def __init__(
-        self, function: typing.Callable[[typing.Any], typing.Any], axis: int = 0
+        self, function: typing.Callable[[typing.Any], typing.Any], **kwargs
     ) -> None:
-        super().__init__(function=lambda a: (a, function(a)), axis=axis)
+        super().__init__(function=lambda a: (a, function(a)), **kwargs)
 
 
 assert range(65, 68) | PipePair(chr) | Pipe(list) == [
@@ -232,3 +259,6 @@ assert range(65, 68) | PipePair(chr) | Pipe(list) == [
 assert range(2, 4) | PipeMap(range) | PipePair(lambda n: n**3, axis=1) | PipeMap(
     set
 ) | Pipe(list) == [{(0, 0), (1, 1)}, {(0, 0), (1, 1), (2, 8)}]
+assert [1, 0, 2] | PipePair(lambda d: 42 // d, catch=True) | PipeMap(
+    lambda r: (r[0], r[1].args) if isinstance(r[1], Exception) else r
+) | Pipe(list) == [(1, 42), (0, ("integer division or modulo by zero",)), (2, 21)]