Hacks to Python | Nan's Blog

Pipeline operator

Models

Miss your dear and dear pipeline operator in …

1
cat file.txt | grep "hello" | wc -l

1
iris %>% filter(Species == "setosa") %>% select(Sepal.Length, Sepal.Width)

1
import Data.Function ((&))
2

3
[1, 2, 3, 4, 5] & filter odd & map (*2) & sum

Well, you can have it in Python too!

Implementation

1
from functools import update_wrapper
2
from typing import Any
3
from collections.abc import Callable
4

5

6
class Pipeline:
7
    def __init__(self, func: Callable | None = None) -> None:
8
        if func is not None:
9
            update_wrapper(self, func)
10
            self.funcs: list[Callable] = [func]
11
        else:
12
            self.funcs: list[Callable] = []
13

14
    def __call__(self, *args: Any, **kwds: Any) -> Any:
15
        """
16
        self(*) => self.funcs[0](self.funcs[1](self.funcs[2](...)))
17
        """
18
        if len(self.funcs) == 0:
19
            return None
20
        res = self.funcs[-1](*args, **kwds)
21
        for func in reversed(self.funcs[:-1]):
22
            try:
23
                res = func(res)
24
            except Exception as e:
25
                print(f"Error encountered when executing {func} with parameter {res}:")
26
                raise e
27
        return res
28

29
    def __add__(self, func: Callable | "Pipeline") -> "Pipeline":
30
        """
31
        (f + g)(*) => f(g(*))
32
        """
33
        res = Pipeline()
34
        res.funcs += self.funcs
35
        if isinstance(func, Pipeline):
36
            res.funcs += func.funcs
37
        elif callable(func):
38
            res.funcs.append(func)
39
        else:
40
            raise TypeError(f"Unsupported type {type(func)}")
41
        return res
42

43
    def __radd__(self, func: Callable | "Pipeline") -> "Pipeline":
44
        res = Pipeline()
45
        if isinstance(func, Pipeline):
46
            res.funcs += func.funcs
47
        elif callable(func):
48
            res.funcs.append(func)
49
        else:
50
            raise TypeError(f"Unsupported type {type(func)}")
51
        res.funcs += self.funcs
52
        return res
53

54
    def __or__(self, func: Callable | "Pipeline"):
55
        """
56
        * | f | g => g(f(*))
57
        """
58
        return func + self
59

60
    def __ror__(self, func: Callable | "Pipeline" | Any):
61
        if isinstance(func, Callable) or isinstance(func, Pipeline):
62
            return self + func
63
        else:
64
            return self(func)
65

66

67
@Pipeline
68
def foo(a: int) -> list[int]:
69
    return list(range(a))
70

71

72
@Pipeline
73
def bar(b: list[int]) -> int:
74
    return max(b)
75

76

77
# Pipeline in Python???!!!
78
if __name__ == "__main__":
79
    func = bar + foo
80
    print(func(10))  # 9
81
    print(10 | foo | bar)  # 9

How it works

The solution is simple:

Define a wrapper function that keeps the original functions inside (self.funcs is a list of callables)
Generate a new function upon operating at least one such function (e.g. pipeline operator | which is seen as or operation in Python, or adding two functions together like f+g <=> f(g(.)))
When calling such function, it will sequentially execute all the nested functions
Add some finishing touches, and you are good to go!

Better type hinting

Check out this project of Yixuan Wang that provides better type hinting for similar pipeline operator implementation.

Multi-statement lambda expression

Models

Multi-statement lambda expressions are commonly supported in other languages, such as:

1
auto func = [](auto a, auto b)
2
{
3
    a++;
4
    b++;
5
    return a + b;
6
};

1
var func = (int a, int b) =>
2
{
3
    a++;
4
    b++;
5
    return a + b;
6
};

1
func := func(a int, b int) int {
2
    a++
3
    b++
4
    return a + b
5
}

1
const func = (a: number, b: number): number => {
2
  a++;
3
  b++;
4
  return a + b;
5
};

1
const func = (a, b) => {
2
  a++;
3
  b++;
4
  return a + b;
5
};

1
let func = |mut a: i32, mut b: i32| -> i32 {
2
    a += 1;
3
    b += 1;
4
    a + b
5
};

But wait! What about Python?

1
func = lambda a, b: a + b

Seems like Python only supports single-statement lambda expression… or is it?

Well, you can have it too!

Implementation

1
func = lambda a, b: (
2
    a := a + 1,
3
    b := b + 1,
4
    a + b,
5
)[-1]

How it works

Well, in python, the evaluation of a tuple is done sequentially, so you can use a tuple to mimic the behavior of a multi-statement lambda expression, where the return value could be specified by the index of element.

Note that in order to assign values in a tuple, you will need to use the walrus operator := introduced in Python 3.8. It will assign the value to the variable on the left, and return the value on the right, similar to the default behavior of = in C/C++ and many other languages.