kpipeline 0.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kpipeline-0.1.1/PKG-INFO +74 -0
- kpipeline-0.1.1/README.md +67 -0
- kpipeline-0.1.1/kpipeline/__init__.py +0 -0
- kpipeline-0.1.1/kpipeline/async_pipeline.py +442 -0
- kpipeline-0.1.1/kpipeline/graph.py +109 -0
- kpipeline-0.1.1/kpipeline/pipeline.py +451 -0
- kpipeline-0.1.1/kpipeline/py.typed +0 -0
- kpipeline-0.1.1/kpipeline.egg-info/PKG-INFO +74 -0
- kpipeline-0.1.1/kpipeline.egg-info/SOURCES.txt +13 -0
- kpipeline-0.1.1/kpipeline.egg-info/dependency_links.txt +1 -0
- kpipeline-0.1.1/kpipeline.egg-info/top_level.txt +1 -0
- kpipeline-0.1.1/pyproject.toml +14 -0
- kpipeline-0.1.1/setup.cfg +4 -0
- kpipeline-0.1.1/test/test_graph.py +67 -0
- kpipeline-0.1.1/test/test_graph_dot.py +35 -0
kpipeline-0.1.1/PKG-INFO
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: kpipeline
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: A pipeline library for Python
|
|
5
|
+
Requires-Python: >=3.13
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
|
|
8
|
+
# kpipe
|
|
9
|
+
|
|
10
|
+
`kpipe` is a simple pipeline library that allows you to write your complex application as a pipeline. This helps code organization, debugging, and testing.
|
|
11
|
+
|
|
12
|
+
## How to use
|
|
13
|
+
|
|
14
|
+
kpipe is like a DSL language inside Python. It is based on "pipes" which can be combined into pipelines using the pipe primitives provided by the library.
|
|
15
|
+
|
|
16
|
+
Each pipe represents a function transforming an input into an output. Pipes are stateless and immutable. Below is an example defining two example pipes.
|
|
17
|
+
|
|
18
|
+
```python
|
|
19
|
+
from kpipeline import Pipe
|
|
20
|
+
|
|
21
|
+
class AddOnePipe(Pipe[int, int, None]):
|
|
22
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
23
|
+
return input + 1
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class MulByTwoPipe(Pipe[int, int, None]):
|
|
27
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
28
|
+
return input * 2
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
These pipes can be combined using the ChainPipe primitive (aliased into the `|` operator) two form a pipeline that performs these two pipes sequentially:
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
pipeline = AddOnePipe() | MulByTwoPipe()
|
|
35
|
+
|
|
36
|
+
# or
|
|
37
|
+
|
|
38
|
+
from kpipeline import ChainPipe
|
|
39
|
+
pipeline = ChainPipe(AddOnePipe(), MulByTwoPipe())
|
|
40
|
+
|
|
41
|
+
print(pipeline.apply(2)) # 6
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
Currently, the library defines these primitives:
|
|
45
|
+
|
|
46
|
+
| Primitive | Purpose |
|
|
47
|
+
|---|---|
|
|
48
|
+
|`ChainPipe`|Execute two pipes sequentially|
|
|
49
|
+
|`ConditionalPipe`|Execute a pipe if a given condition is true|
|
|
50
|
+
|`BranchPipe`|Execute one of two pipes depending on whether the condition is true or false|
|
|
51
|
+
|`SelectPipe`|Execute one of multiple pipes based on a selector key|
|
|
52
|
+
|`ParallelPipe`|Execute multiple pipes and combine their results (only async implementation is actually parallel)|
|
|
53
|
+
|`MetadataWrapperPipe`|Transform the metadata given to the pipeline for the subpipe|
|
|
54
|
+
|`MapPipe`|Apply a subpipe into a sequence of inputs|
|
|
55
|
+
|`FilterPipe`|Filter a sequence of inputs using a predicate|
|
|
56
|
+
|`RetryPipe`|Run a pipe multiple times in case it fails|
|
|
57
|
+
|
|
58
|
+
## Tests
|
|
59
|
+
|
|
60
|
+
You can run all tests in this repository with
|
|
61
|
+
|
|
62
|
+
```
|
|
63
|
+
uv run pytest test
|
|
64
|
+
uv run mypy .
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## License
|
|
68
|
+
|
|
69
|
+
TBD
|
|
70
|
+
|
|
71
|
+
## AI Use Disclosure
|
|
72
|
+
|
|
73
|
+
I hate writing tests so unit tests have been generated by GPT-OSS 120B for your convenience.
|
|
74
|
+
They are provided in the hope that they are better than nothing.
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
# kpipe
|
|
2
|
+
|
|
3
|
+
`kpipe` is a simple pipeline library that allows you to write your complex application as a pipeline. This helps code organization, debugging, and testing.
|
|
4
|
+
|
|
5
|
+
## How to use
|
|
6
|
+
|
|
7
|
+
kpipe is like a DSL language inside Python. It is based on "pipes" which can be combined into pipelines using the pipe primitives provided by the library.
|
|
8
|
+
|
|
9
|
+
Each pipe represents a function transforming an input into an output. Pipes are stateless and immutable. Below is an example defining two example pipes.
|
|
10
|
+
|
|
11
|
+
```python
|
|
12
|
+
from kpipeline import Pipe
|
|
13
|
+
|
|
14
|
+
class AddOnePipe(Pipe[int, int, None]):
|
|
15
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
16
|
+
return input + 1
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class MulByTwoPipe(Pipe[int, int, None]):
|
|
20
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
21
|
+
return input * 2
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
These pipes can be combined using the ChainPipe primitive (aliased into the `|` operator) two form a pipeline that performs these two pipes sequentially:
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
pipeline = AddOnePipe() | MulByTwoPipe()
|
|
28
|
+
|
|
29
|
+
# or
|
|
30
|
+
|
|
31
|
+
from kpipeline import ChainPipe
|
|
32
|
+
pipeline = ChainPipe(AddOnePipe(), MulByTwoPipe())
|
|
33
|
+
|
|
34
|
+
print(pipeline.apply(2)) # 6
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
Currently, the library defines these primitives:
|
|
38
|
+
|
|
39
|
+
| Primitive | Purpose |
|
|
40
|
+
|---|---|
|
|
41
|
+
|`ChainPipe`|Execute two pipes sequentially|
|
|
42
|
+
|`ConditionalPipe`|Execute a pipe if a given condition is true|
|
|
43
|
+
|`BranchPipe`|Execute one of two pipes depending on whether the condition is true or false|
|
|
44
|
+
|`SelectPipe`|Execute one of multiple pipes based on a selector key|
|
|
45
|
+
|`ParallelPipe`|Execute multiple pipes and combine their results (only async implementation is actually parallel)|
|
|
46
|
+
|`MetadataWrapperPipe`|Transform the metadata given to the pipeline for the subpipe|
|
|
47
|
+
|`MapPipe`|Apply a subpipe into a sequence of inputs|
|
|
48
|
+
|`FilterPipe`|Filter a sequence of inputs using a predicate|
|
|
49
|
+
|`RetryPipe`|Run a pipe multiple times in case it fails|
|
|
50
|
+
|
|
51
|
+
## Tests
|
|
52
|
+
|
|
53
|
+
You can run all tests in this repository with
|
|
54
|
+
|
|
55
|
+
```
|
|
56
|
+
uv run pytest test
|
|
57
|
+
uv run mypy .
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
## License
|
|
61
|
+
|
|
62
|
+
TBD
|
|
63
|
+
|
|
64
|
+
## AI Use Disclosure
|
|
65
|
+
|
|
66
|
+
I hate writing tests so unit tests have been generated by GPT-OSS 120B for your convenience.
|
|
67
|
+
They are provided in the hope that they are better than nothing.
|
|
File without changes
|
|
@@ -0,0 +1,442 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
from collections.abc import Awaitable, Mapping, Sequence
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from typing import Callable, Optional
|
|
5
|
+
|
|
6
|
+
from .graph import Graph, GraphNode, GraphConnection
|
|
7
|
+
from .pipeline import BasePipe, Pipe
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class AsyncPipe[Input, Output, Metadata](BasePipe[Input, Awaitable[Output], Metadata]):
|
|
11
|
+
"""
|
|
12
|
+
Abstract base class of all asynchronous pipes.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
16
|
+
"""
|
|
17
|
+
Like batch_apply, but returns Awaitable[Sequence[Output]] instead of Sequence[Awaitable[Output]].
|
|
18
|
+
The default implementation uses asyncio.gather.
|
|
19
|
+
"""
|
|
20
|
+
return await asyncio.gather(*self.batch_apply(data, metadata))
|
|
21
|
+
|
|
22
|
+
def async_chain[OtherOutput](self, other: "Pipe[Output, OtherOutput, Metadata] | AsyncPipe[Output, OtherOutput, Metadata]") -> "AsyncChainPipe[Input, Output, OtherOutput, Metadata]":
|
|
23
|
+
"""
|
|
24
|
+
Apply the second pipe to the result of the first pipe. Same as AsyncChainPipe(self, other).
|
|
25
|
+
"""
|
|
26
|
+
return AsyncChainPipe(self, other)
|
|
27
|
+
|
|
28
|
+
def __or__[OtherOutput](self, other: "Pipe[Output, OtherOutput, Metadata] | AsyncPipe[Output, OtherOutput, Metadata]") -> "AsyncChainPipe[Input, Output, OtherOutput, Metadata]":
|
|
29
|
+
"""
|
|
30
|
+
Alias for AsyncPipe.async_chain.
|
|
31
|
+
"""
|
|
32
|
+
return self.async_chain(other)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
type SyncOrAsyncPipe[I, O, M] = Pipe[I, O, M] | AsyncPipe[I, O, M]
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
async def _await_or_return[T](data: T | Awaitable[T]) -> T:
|
|
39
|
+
if isinstance(data, Awaitable):
|
|
40
|
+
return await data
|
|
41
|
+
|
|
42
|
+
else:
|
|
43
|
+
return data
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
type SyncOrAsyncPipeOrCallable[I, O, M] = Callable[[I, M], O] | Callable[[I, M], Awaitable[O]] | SyncOrAsyncPipe[I, O, M]
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
async def _call_or_apply[I, O, M](f: SyncOrAsyncPipeOrCallable[I, O, M], d: I, m: M) -> O:
|
|
50
|
+
if isinstance(f, BasePipe):
|
|
51
|
+
return await _await_or_return(f.apply(d, m))
|
|
52
|
+
|
|
53
|
+
else:
|
|
54
|
+
return await _await_or_return(f(d, m))
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
async def _batch_apply[I, O, M](p: SyncOrAsyncPipe[I, O, M], d: Sequence[I], m: M) -> Sequence[O]:
|
|
58
|
+
if isinstance(p, AsyncPipe):
|
|
59
|
+
return await p.async_batch_apply(d, m)
|
|
60
|
+
|
|
61
|
+
else:
|
|
62
|
+
return p.batch_apply(d, m)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
@dataclass(frozen=True)
|
|
66
|
+
class AsyncChainPipe[Input, Middle, Output, Metadata](AsyncPipe[Input, Output, Metadata]):
|
|
67
|
+
"""
|
|
68
|
+
Connects to pipes so that the second pipe is applied to the result of the first pipe.
|
|
69
|
+
"""
|
|
70
|
+
pipe1: SyncOrAsyncPipe[Input, Middle, Metadata]
|
|
71
|
+
pipe2: SyncOrAsyncPipe[Middle, Output, Metadata]
|
|
72
|
+
|
|
73
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
74
|
+
middle = await _await_or_return(self.pipe1.apply(data, metadata))
|
|
75
|
+
return await _await_or_return(self.pipe2.apply(middle, metadata))
|
|
76
|
+
|
|
77
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
78
|
+
middle_batch = await _batch_apply(self.pipe1, data, metadata)
|
|
79
|
+
return await _batch_apply(self.pipe2, middle_batch, metadata)
|
|
80
|
+
|
|
81
|
+
def to_graph(self) -> Graph:
|
|
82
|
+
graph1 = self.pipe1.to_graph()
|
|
83
|
+
graph2 = self.pipe2.to_graph()
|
|
84
|
+
return graph1 >> graph2
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
@dataclass(frozen=True)
|
|
88
|
+
class AsyncIdentityPipe[InputOutput, Metadata](AsyncPipe[InputOutput, InputOutput, Metadata]):
|
|
89
|
+
"""
|
|
90
|
+
A pipe that returns its input.
|
|
91
|
+
"""
|
|
92
|
+
async def apply(self, data: InputOutput, metadata: Metadata) -> InputOutput:
|
|
93
|
+
return data
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
@dataclass(frozen=True)
|
|
97
|
+
class AsyncMerge2Pipe[Input, Output1, Output2, Output, Metadata](AsyncPipe[Input, Output, Metadata]):
|
|
98
|
+
"""
|
|
99
|
+
Runs two pipes and merges their results.
|
|
100
|
+
"""
|
|
101
|
+
pipe1: SyncOrAsyncPipe[Input, Output1, Metadata]
|
|
102
|
+
pipe2: SyncOrAsyncPipe[Input, Output2, Metadata]
|
|
103
|
+
merge: Callable[[Output1, Output2], Output]
|
|
104
|
+
description: str = "Combine results"
|
|
105
|
+
|
|
106
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
107
|
+
output1 = await _await_or_return(self.pipe1.apply(data, metadata))
|
|
108
|
+
output2 = await _await_or_return(self.pipe2.apply(data, metadata))
|
|
109
|
+
return self.merge(output1, output2)
|
|
110
|
+
|
|
111
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
112
|
+
batch1 = await _batch_apply(self.pipe1, data, metadata)
|
|
113
|
+
batch2 = await _batch_apply(self.pipe2, data, metadata)
|
|
114
|
+
|
|
115
|
+
assert len(data) == len(batch1) == len(batch2), "lengths do not match"
|
|
116
|
+
return await asyncio.gather(*[_await_or_return(self.merge(output1, output2)) for output1, output2 in zip(batch1, batch2)])
|
|
117
|
+
|
|
118
|
+
def to_graph(self) -> Graph:
|
|
119
|
+
graph1 = self.pipe1.to_graph()
|
|
120
|
+
graph2 = self.pipe2.to_graph()
|
|
121
|
+
combine_node = self.to_node()._replace(title=self.description, shape="combine")
|
|
122
|
+
graph = Graph(nodes=(combine_node,), inputs=(combine_node.id,), outputs=(combine_node.id,))
|
|
123
|
+
return (graph1 | graph2) >> graph
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
@dataclass(frozen=True)
|
|
127
|
+
class AsyncMerge3Pipe[Input, Output1, Output2, Output3, Output, Metadata](AsyncPipe[Input, Output, Metadata]):
|
|
128
|
+
"""
|
|
129
|
+
Runs three pipes and merges their results.
|
|
130
|
+
"""
|
|
131
|
+
pipe1: SyncOrAsyncPipe[Input, Output1, Metadata]
|
|
132
|
+
pipe2: SyncOrAsyncPipe[Input, Output2, Metadata]
|
|
133
|
+
pipe3: SyncOrAsyncPipe[Input, Output3, Metadata]
|
|
134
|
+
merge: Callable[[Output1, Output2, Output3], Output] | Callable[[Output1, Output2, Output3], Awaitable[Output]]
|
|
135
|
+
description: str = "Combine results"
|
|
136
|
+
|
|
137
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
138
|
+
output1 = await _await_or_return(self.pipe1.apply(data, metadata))
|
|
139
|
+
output2 = await _await_or_return(self.pipe2.apply(data, metadata))
|
|
140
|
+
output3 = await _await_or_return(self.pipe3.apply(data, metadata))
|
|
141
|
+
return await _await_or_return(self.merge(output1, output2, output3))
|
|
142
|
+
|
|
143
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
144
|
+
batch1 = await _batch_apply(self.pipe1, data, metadata)
|
|
145
|
+
batch2 = await _batch_apply(self.pipe2, data, metadata)
|
|
146
|
+
batch3 = await _batch_apply(self.pipe3, data, metadata)
|
|
147
|
+
|
|
148
|
+
assert len(data) == len(batch1) == len(batch2) == len(batch3), "lengths do not match"
|
|
149
|
+
return await asyncio.gather(*[_await_or_return(self.merge(output1, output2, output3)) for output1, output2, output3 in zip(batch1, batch2, batch3)])
|
|
150
|
+
|
|
151
|
+
def to_graph(self) -> Graph:
|
|
152
|
+
graph1 = self.pipe1.to_graph()
|
|
153
|
+
graph2 = self.pipe2.to_graph()
|
|
154
|
+
graph3 = self.pipe3.to_graph()
|
|
155
|
+
combine_node = self.to_node()._replace(title=self.description, shape="combine")
|
|
156
|
+
graph = Graph(nodes=(combine_node,), inputs=(combine_node.id,), outputs=(combine_node.id,))
|
|
157
|
+
return (graph1 | graph2 | graph3) >> graph
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
@dataclass(frozen=True)
|
|
161
|
+
class AsyncBranchPipe[Input, Output, Metadata](AsyncPipe[Input, Output, Metadata]):
|
|
162
|
+
"""
|
|
163
|
+
Selects one of two pipes to apply based on a condition.
|
|
164
|
+
"""
|
|
165
|
+
condition: SyncOrAsyncPipeOrCallable[Input, bool, Metadata]
|
|
166
|
+
then_pipe: SyncOrAsyncPipe[Input, Output, Metadata]
|
|
167
|
+
else_pipe: SyncOrAsyncPipe[Input, Output, Metadata]
|
|
168
|
+
description: str = ""
|
|
169
|
+
|
|
170
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
171
|
+
if await _call_or_apply(self.condition, data, metadata):
|
|
172
|
+
return await _await_or_return(self.then_pipe.apply(data, metadata))
|
|
173
|
+
|
|
174
|
+
else:
|
|
175
|
+
return await _await_or_return(self.else_pipe.apply(data, metadata))
|
|
176
|
+
|
|
177
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
178
|
+
conds = await asyncio.gather(*[_call_or_apply(self.condition, d, metadata) for d in data])
|
|
179
|
+
then_batch_future = _batch_apply(self.then_pipe, [d for c, d in zip(conds, data) if c], metadata)
|
|
180
|
+
else_batch_future = _batch_apply(self.else_pipe, [d for c, d in zip(conds, data) if not c], metadata)
|
|
181
|
+
then_batch, else_batch = await asyncio.gather(then_batch_future, else_batch_future)
|
|
182
|
+
|
|
183
|
+
ans: list[Output] = []
|
|
184
|
+
t = 0
|
|
185
|
+
e = 0
|
|
186
|
+
for cond in conds:
|
|
187
|
+
if cond:
|
|
188
|
+
ans.append(then_batch[t])
|
|
189
|
+
t += 1
|
|
190
|
+
|
|
191
|
+
else:
|
|
192
|
+
ans.append(else_batch[e])
|
|
193
|
+
e += 1
|
|
194
|
+
|
|
195
|
+
return ans
|
|
196
|
+
|
|
197
|
+
def to_graph(self) -> Graph:
|
|
198
|
+
condition_node = self.to_node()._replace(title=self.description or "Condition", shape="condition", subgraph=self.condition.to_graph() if isinstance(self.condition, Pipe) else None)
|
|
199
|
+
then_graph = self.then_pipe.to_graph()
|
|
200
|
+
else_graph = self.else_pipe.to_graph()
|
|
201
|
+
return (
|
|
202
|
+
(then_graph | else_graph)
|
|
203
|
+
.add(nodes=(condition_node,))
|
|
204
|
+
.add(connections=tuple(GraphConnection(condition_node.id, i, "yes") for i in then_graph.inputs))
|
|
205
|
+
.add(connections=tuple(GraphConnection(condition_node.id, i, "no") for i in else_graph.inputs))
|
|
206
|
+
._replace(
|
|
207
|
+
inputs=(condition_node.id,),
|
|
208
|
+
outputs=then_graph.outputs + else_graph.outputs,
|
|
209
|
+
)
|
|
210
|
+
)
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
@dataclass(frozen=True)
|
|
214
|
+
class AsyncConditionalPipe[InputOutput, Metadata](AsyncPipe[InputOutput, InputOutput, Metadata]):
|
|
215
|
+
"""
|
|
216
|
+
Applies the pipe only if the condition is met. Similar to AsyncBranchPipe except that the else branch is an identity function.
|
|
217
|
+
Can only be used when the subpipe returns the same type that gets in.
|
|
218
|
+
"""
|
|
219
|
+
condition: Callable[[InputOutput, Metadata], bool | Awaitable[bool]]
|
|
220
|
+
subpipe: SyncOrAsyncPipe[InputOutput, InputOutput, Metadata]
|
|
221
|
+
description: str = ""
|
|
222
|
+
|
|
223
|
+
async def apply(self, data: InputOutput, metadata: Metadata) -> InputOutput:
|
|
224
|
+
if await _await_or_return(self.condition(data, metadata)):
|
|
225
|
+
return await _await_or_return(self.subpipe.apply(data, metadata))
|
|
226
|
+
|
|
227
|
+
else:
|
|
228
|
+
return data
|
|
229
|
+
|
|
230
|
+
async def async_batch_apply(self, data: Sequence[InputOutput], metadata: Metadata) -> Sequence[InputOutput]:
|
|
231
|
+
conds = await asyncio.gather(*[_call_or_apply(self.condition, d, metadata) for d in data])
|
|
232
|
+
then_batch = await _batch_apply(self.subpipe, [d for c, d in zip(conds, data) if c], metadata)
|
|
233
|
+
|
|
234
|
+
ans: list[InputOutput] = []
|
|
235
|
+
t = 0
|
|
236
|
+
for i, cond in enumerate(conds):
|
|
237
|
+
if cond:
|
|
238
|
+
ans.append(then_batch[t])
|
|
239
|
+
t += 1
|
|
240
|
+
|
|
241
|
+
else:
|
|
242
|
+
ans.append(data[i])
|
|
243
|
+
|
|
244
|
+
return ans
|
|
245
|
+
|
|
246
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
247
|
+
return self.subpipe.to_graph()
|
|
248
|
+
|
|
249
|
+
def to_node(self) -> GraphNode:
|
|
250
|
+
return super().to_node()._replace(title=self.description)
|
|
251
|
+
|
|
252
|
+
def is_wrapper(self) -> bool:
|
|
253
|
+
return True
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
@dataclass(frozen=True)
|
|
257
|
+
class AsyncSelectPipe[Input, Output, Metadata, Key](AsyncPipe[Input, Output, Metadata]):
|
|
258
|
+
"""
|
|
259
|
+
Selects one of multiple pipes based on a selector function. AsyncBranchPipe generalized to more than two branches.
|
|
260
|
+
"""
|
|
261
|
+
key: SyncOrAsyncPipeOrCallable[Input, Key, Metadata]
|
|
262
|
+
subpipes: Mapping[Key, SyncOrAsyncPipe[Input, Output, Metadata]]
|
|
263
|
+
otherwise_pipe: SyncOrAsyncPipe[Input, Output, Metadata]
|
|
264
|
+
description: str = ""
|
|
265
|
+
|
|
266
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
267
|
+
key = await _call_or_apply(self.key, data, metadata)
|
|
268
|
+
if key in self.subpipes:
|
|
269
|
+
return await _await_or_return(self.subpipes[key].apply(data, metadata))
|
|
270
|
+
|
|
271
|
+
else:
|
|
272
|
+
return await _await_or_return(self.otherwise_pipe.apply(data, metadata))
|
|
273
|
+
|
|
274
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
275
|
+
keys = await asyncio.gather(*[_call_or_apply(self.key, d, metadata) for d in data])
|
|
276
|
+
unique_keys = list(set(keys))
|
|
277
|
+
batches: dict[Key, list[tuple[int, Input]]] = {key: [] for key in unique_keys}
|
|
278
|
+
for i, (key, item) in enumerate(zip(keys, data)):
|
|
279
|
+
batches[key].append((i, item))
|
|
280
|
+
|
|
281
|
+
batch_futures: list[Awaitable[Sequence[Output]]] = []
|
|
282
|
+
for key in unique_keys:
|
|
283
|
+
pipe = self.subpipes.get(key, self.otherwise_pipe)
|
|
284
|
+
batch = [item for _, item in batches[key]]
|
|
285
|
+
batch_futures.append(_batch_apply(pipe, batch, metadata))
|
|
286
|
+
|
|
287
|
+
batch_results = await asyncio.gather(*batch_futures)
|
|
288
|
+
batch_results_dict = {key: batch_results[i] for i, key in enumerate(unique_keys)}
|
|
289
|
+
ans_dict = {i: result for key, batch in batches.items() for (i, _), result in zip(batch, batch_results_dict[key])}
|
|
290
|
+
return [ans_dict[i] for i in range(len(data))]
|
|
291
|
+
|
|
292
|
+
def to_graph(self) -> Graph:
|
|
293
|
+
condition_node = self.to_node()._replace(title=self.description or "Condition", shape="condition", subgraph=self.key.to_graph() if isinstance(self.key, Pipe) else None)
|
|
294
|
+
otherwise_graph = self.otherwise_pipe.to_graph()
|
|
295
|
+
graph = otherwise_graph.add(
|
|
296
|
+
nodes=(condition_node,),
|
|
297
|
+
connections=tuple(GraphConnection(condition_node.id, i, "otherwise") for i in otherwise_graph.inputs),
|
|
298
|
+
)
|
|
299
|
+
for key, subpipe in self.subpipes.items():
|
|
300
|
+
subpipe_graph = subpipe.to_graph()
|
|
301
|
+
graph |= subpipe_graph
|
|
302
|
+
graph = graph.add(connections=tuple(GraphConnection(condition_node.id, i, str(key)) for i in subpipe_graph.inputs))
|
|
303
|
+
|
|
304
|
+
return graph._replace(inputs=(condition_node.id,))
|
|
305
|
+
|
|
306
|
+
|
|
307
|
+
@dataclass(frozen=True)
|
|
308
|
+
class AsyncParallelPipe[Input, Output, CombinedOutput, Metadata](AsyncPipe[Input, CombinedOutput, Metadata]):
|
|
309
|
+
"""
|
|
310
|
+
Executes multiple pipes in parallel on the same input and combines the outputs.
|
|
311
|
+
"""
|
|
312
|
+
subpipes: Sequence[SyncOrAsyncPipe[Input, Output, Metadata]]
|
|
313
|
+
combine: SyncOrAsyncPipeOrCallable[Sequence[Output], CombinedOutput, Metadata]
|
|
314
|
+
description: str = ""
|
|
315
|
+
|
|
316
|
+
async def apply(self, data: Input, metadata: Metadata) -> CombinedOutput:
|
|
317
|
+
results: list[Awaitable[Output]] = []
|
|
318
|
+
for subpipe in self.subpipes:
|
|
319
|
+
results.append(_await_or_return(subpipe.apply(data, metadata)))
|
|
320
|
+
|
|
321
|
+
gathered: Sequence[Output] = await asyncio.gather(*results)
|
|
322
|
+
return await _call_or_apply(self.combine, gathered, metadata)
|
|
323
|
+
|
|
324
|
+
async def async_batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[CombinedOutput]:
|
|
325
|
+
results: list[Awaitable[Sequence[Output]]] = []
|
|
326
|
+
for subpipe in self.subpipes:
|
|
327
|
+
results.append(_batch_apply(subpipe, data, metadata))
|
|
328
|
+
|
|
329
|
+
gathered = await asyncio.gather(*results)
|
|
330
|
+
zipped_results: zip[Sequence[Output]] = zip(*gathered)
|
|
331
|
+
return await asyncio.gather(*[_call_or_apply(self.combine, seq, metadata) for seq in zipped_results])
|
|
332
|
+
|
|
333
|
+
def to_graph(self) -> Graph:
|
|
334
|
+
combine_node = self.to_node()._replace(title=self.description or "Combine results", shape="combine", subgraph=self.combine.to_graph() if isinstance(self.combine, Pipe) else None)
|
|
335
|
+
graph = Graph(nodes=(combine_node,))
|
|
336
|
+
for subpipe in self.subpipes:
|
|
337
|
+
subpipe_graph = subpipe.to_graph()
|
|
338
|
+
graph |= subpipe_graph
|
|
339
|
+
graph = graph.add(connections=tuple(GraphConnection(o, combine_node.id) for o in subpipe_graph.outputs))
|
|
340
|
+
|
|
341
|
+
return graph._replace(outputs=(combine_node.id,))
|
|
342
|
+
|
|
343
|
+
|
|
344
|
+
@dataclass(frozen=True)
|
|
345
|
+
class AsyncMetadataWrapperPipe[Input, Output, OuterMetadata, InnerMetadata](AsyncPipe[Input, Output, OuterMetadata]):
|
|
346
|
+
"""
|
|
347
|
+
Executes a pipe with changed metadata.
|
|
348
|
+
|
|
349
|
+
Does NOT allow to_inner to be a pipe: this is intentional, metadata is not transformed by pipes.
|
|
350
|
+
"""
|
|
351
|
+
to_inner: Callable[[OuterMetadata], InnerMetadata | Awaitable[InnerMetadata]]
|
|
352
|
+
subpipe: SyncOrAsyncPipe[Input, Output, InnerMetadata]
|
|
353
|
+
description: str = ""
|
|
354
|
+
|
|
355
|
+
async def apply(self, data: Input, metadata: OuterMetadata) -> Output:
|
|
356
|
+
return await _await_or_return(self.subpipe.apply(data, await _await_or_return(self.to_inner(metadata))))
|
|
357
|
+
|
|
358
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
359
|
+
return self.subpipe.to_graph()
|
|
360
|
+
|
|
361
|
+
def to_node(self) -> GraphNode:
|
|
362
|
+
return super().to_node()._replace(title=self.description)
|
|
363
|
+
|
|
364
|
+
def is_wrapper(self) -> bool:
|
|
365
|
+
return True
|
|
366
|
+
|
|
367
|
+
|
|
368
|
+
@dataclass(frozen=True)
|
|
369
|
+
class AsyncMapPipe[Input, Output, Metadata](AsyncPipe[Sequence[Input], Sequence[Output], Metadata]):
|
|
370
|
+
"""
|
|
371
|
+
Maps a sequence of input objects into a sequence of output objects with a subpipe.
|
|
372
|
+
|
|
373
|
+
Does NOT allow the subpipe to be a callable: this is intentional, functions that transforms Inputs to Outputs should be pipes.
|
|
374
|
+
"""
|
|
375
|
+
subpipe: SyncOrAsyncPipe[Input, Output, Metadata]
|
|
376
|
+
description: str = "Map"
|
|
377
|
+
|
|
378
|
+
async def apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
379
|
+
if isinstance(self.subpipe, AsyncPipe):
|
|
380
|
+
return await self.subpipe.async_batch_apply(data, metadata)
|
|
381
|
+
|
|
382
|
+
else:
|
|
383
|
+
return self.subpipe.batch_apply(data, metadata)
|
|
384
|
+
|
|
385
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
386
|
+
return self.subpipe.to_graph()
|
|
387
|
+
|
|
388
|
+
def to_node(self) -> GraphNode:
|
|
389
|
+
return super().to_node()._replace(title=self.description)
|
|
390
|
+
|
|
391
|
+
|
|
392
|
+
@dataclass(frozen=True)
|
|
393
|
+
class AsyncFilterPipe[Input, Metadata](AsyncPipe[Sequence[Input], Sequence[Input], Metadata]):
|
|
394
|
+
"""
|
|
395
|
+
Filters a sequence of input object with a predicate.
|
|
396
|
+
"""
|
|
397
|
+
predicate: SyncOrAsyncPipeOrCallable[Input, bool, Metadata]
|
|
398
|
+
description: str = "Filter"
|
|
399
|
+
|
|
400
|
+
async def apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Input]:
|
|
401
|
+
return [i for i in data if await _call_or_apply(self.predicate, i, metadata)]
|
|
402
|
+
|
|
403
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
404
|
+
return self.predicate.to_graph() if isinstance(self.predicate, BasePipe) else None
|
|
405
|
+
|
|
406
|
+
def to_node(self) -> GraphNode:
|
|
407
|
+
return super().to_node()._replace(title=self.description)
|
|
408
|
+
|
|
409
|
+
|
|
410
|
+
@dataclass(frozen=True)
|
|
411
|
+
class AsyncRetryPipe[Input, Output, Metadata](AsyncPipe[Input, Output, Metadata]):
|
|
412
|
+
"""
|
|
413
|
+
Tries to execute the subpipe several times in case it fails.
|
|
414
|
+
"""
|
|
415
|
+
subpipe: SyncOrAsyncPipe[Input, Output, Metadata]
|
|
416
|
+
retries: int
|
|
417
|
+
exceptions: type | tuple[type, ...]
|
|
418
|
+
description: str = "Retry several times"
|
|
419
|
+
|
|
420
|
+
async def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
421
|
+
attempt = 0
|
|
422
|
+
while True:
|
|
423
|
+
try:
|
|
424
|
+
return await _await_or_return(self.subpipe.apply(data, metadata))
|
|
425
|
+
|
|
426
|
+
except Exception as e:
|
|
427
|
+
if not isinstance(e, self.exceptions):
|
|
428
|
+
raise e
|
|
429
|
+
|
|
430
|
+
attempt += 1
|
|
431
|
+
if attempt > self.retries:
|
|
432
|
+
raise e
|
|
433
|
+
|
|
434
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
435
|
+
return self.subpipe.to_graph()
|
|
436
|
+
|
|
437
|
+
def to_node(self) -> GraphNode:
|
|
438
|
+
return super().to_node()._replace(title=self.description)
|
|
439
|
+
|
|
440
|
+
def is_wrapper(self) -> bool:
|
|
441
|
+
return True
|
|
442
|
+
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
from typing import Literal, NamedTuple, Optional, Type
|
|
2
|
+
|
|
3
|
+
type GraphNodeId = str
|
|
4
|
+
type GraphNodeShape = Literal["default", "combine", "condition"]
|
|
5
|
+
"""
|
|
6
|
+
The shape of the node.
|
|
7
|
+
`combine` is used for nodes that combine multiple inputs into one output.
|
|
8
|
+
`condition` is used for nodes that diverge path based on condition.
|
|
9
|
+
The actual shape depends on the rendering backend.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class GraphNode(NamedTuple):
|
|
14
|
+
id: GraphNodeId
|
|
15
|
+
title: str
|
|
16
|
+
type: Optional[Type] = None
|
|
17
|
+
shape: GraphNodeShape = "default"
|
|
18
|
+
subgraph: "Optional[Graph]" = None
|
|
19
|
+
|
|
20
|
+
def to_dot(self, indent: int = 0) -> str:
|
|
21
|
+
shape_map = {
|
|
22
|
+
"default": "box",
|
|
23
|
+
"combine": "ellipse",
|
|
24
|
+
"condition": "diamond",
|
|
25
|
+
}
|
|
26
|
+
dot_shape = shape_map.get(self.shape, "box")
|
|
27
|
+
node_def = f"""{" "*indent}"{self.id}" [label="{self.title}", shape={dot_shape}];"""
|
|
28
|
+
if self.subgraph:
|
|
29
|
+
sub_dot = self.subgraph.to_dot(name=None, indent=indent+1)
|
|
30
|
+
node_def += f"""{" "*indent}subgraph "cluster_{self.id}" {{"""
|
|
31
|
+
node_def += f"""{" "*indent} label="{self.title}";"""
|
|
32
|
+
node_def += sub_dot
|
|
33
|
+
node_def += " "*indent + "}"
|
|
34
|
+
return node_def
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class GraphConnection(NamedTuple):
|
|
38
|
+
from_node: GraphNodeId
|
|
39
|
+
to_node: GraphNodeId
|
|
40
|
+
label: str = ""
|
|
41
|
+
|
|
42
|
+
def to_dot(self, indent: int = 0) -> str:
|
|
43
|
+
label_part = f' [label="{self.label}"]' if self.label else ""
|
|
44
|
+
return f"""{" "*indent}"{self.from_node}" -> "{self.to_node}"{label_part};"""
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class Graph(NamedTuple):
|
|
48
|
+
nodes: tuple[GraphNode, ...] = ()
|
|
49
|
+
connections: tuple[GraphConnection, ...] = ()
|
|
50
|
+
inputs: tuple[GraphNodeId, ...] = ()
|
|
51
|
+
outputs: tuple[GraphNodeId, ...] = ()
|
|
52
|
+
|
|
53
|
+
def __or__(self, other: "Graph") -> "Graph":
|
|
54
|
+
return Graph(
|
|
55
|
+
nodes=self.nodes + other.nodes,
|
|
56
|
+
connections=self.connections + other.connections,
|
|
57
|
+
inputs=self.inputs + other.inputs,
|
|
58
|
+
outputs=self.outputs + other.outputs,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
def __rshift__(self, other: "Graph") -> "Graph":
|
|
62
|
+
return Graph(
|
|
63
|
+
nodes=self.nodes + other.nodes,
|
|
64
|
+
connections=self.connections + other.connections + tuple(GraphConnection(o, i) for o in self.outputs for i in other.inputs),
|
|
65
|
+
inputs=self.inputs,
|
|
66
|
+
outputs=other.outputs,
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
def add(self, nodes: tuple[GraphNode, ...] = (), connections: tuple[GraphConnection, ...] = (), inputs: tuple[GraphNodeId, ...] = (), outputs: tuple[GraphNodeId, ...] = ()) -> "Graph":
|
|
70
|
+
return (self | Graph(nodes, connections, inputs, outputs))
|
|
71
|
+
|
|
72
|
+
def is_valid(self):
|
|
73
|
+
"""
|
|
74
|
+
Checks that all connections, inputs, and outputs only refer existing nodes.
|
|
75
|
+
"""
|
|
76
|
+
node_id_set = set()
|
|
77
|
+
node_stack = list(self.nodes)
|
|
78
|
+
while len(node_stack):
|
|
79
|
+
node = node_stack.pop()
|
|
80
|
+
if node.id in node_id_set:
|
|
81
|
+
continue
|
|
82
|
+
|
|
83
|
+
if node.subgraph is not None:
|
|
84
|
+
node_stack += node.subgraph.nodes
|
|
85
|
+
|
|
86
|
+
node_id_set.add(node.id)
|
|
87
|
+
|
|
88
|
+
references = set()
|
|
89
|
+
references |= {c.from_node for c in self.connections}
|
|
90
|
+
references |= {c.to_node for c in self.connections}
|
|
91
|
+
references |= set(self.inputs)
|
|
92
|
+
references |= set(self.outputs)
|
|
93
|
+
return references <= node_id_set
|
|
94
|
+
|
|
95
|
+
def to_dot(self, name: str | None = "graph", indent: int = 0) -> str:
|
|
96
|
+
lines: list[str] = []
|
|
97
|
+
if name is not None:
|
|
98
|
+
lines.append(f"digraph {name} {{")
|
|
99
|
+
|
|
100
|
+
for node in self.nodes:
|
|
101
|
+
lines.append(node.to_dot(indent=indent+1))
|
|
102
|
+
|
|
103
|
+
for conn in self.connections:
|
|
104
|
+
lines.append(conn.to_dot(indent=indent+1))
|
|
105
|
+
|
|
106
|
+
if name is not None:
|
|
107
|
+
lines.append("}")
|
|
108
|
+
|
|
109
|
+
return "\n".join(lines)
|
|
@@ -0,0 +1,451 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
from collections.abc import Mapping, Sequence
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from typing import Callable, Iterable, NamedTuple, Optional
|
|
5
|
+
|
|
6
|
+
from .graph import Graph, GraphNode, GraphConnection
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class BasePipe[Input, Output, Metadata](ABC):
|
|
10
|
+
"""
|
|
11
|
+
Abstract base class of all synchronous and asynchronous pipes.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
@abstractmethod
|
|
15
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
16
|
+
"""
|
|
17
|
+
Apply the pipe to the given input data and return the result.
|
|
18
|
+
"""
|
|
19
|
+
...
|
|
20
|
+
|
|
21
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
22
|
+
"""
|
|
23
|
+
Applies the pipe to multiple data points. By default merely calls apply in sequence.
|
|
24
|
+
"""
|
|
25
|
+
return [self.apply(d, metadata) for d in data]
|
|
26
|
+
|
|
27
|
+
def chain[OtherOutput](self, other: "Pipe[Output, OtherOutput, Metadata]") -> "ChainPipe[Input, Output, OtherOutput, Metadata]":
|
|
28
|
+
"""
|
|
29
|
+
Apply the second pipe to the result of the first pipe. Same as ChainPipe(self, other).
|
|
30
|
+
"""
|
|
31
|
+
return ChainPipe(self, other)
|
|
32
|
+
|
|
33
|
+
def to_node(self) -> GraphNode:
|
|
34
|
+
return GraphNode(id=str(id(self)), type=type(self), title="Pipe")
|
|
35
|
+
|
|
36
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
37
|
+
return None
|
|
38
|
+
|
|
39
|
+
def is_wrapper(self) -> bool:
|
|
40
|
+
return False
|
|
41
|
+
|
|
42
|
+
def to_graph(self) -> Graph:
|
|
43
|
+
"""
|
|
44
|
+
Return the pipeline as a graph.
|
|
45
|
+
"""
|
|
46
|
+
if subgraph := self.get_subgraph():
|
|
47
|
+
node = self.to_node()._replace(subgraph=subgraph)
|
|
48
|
+
return Graph(
|
|
49
|
+
nodes=(node,),
|
|
50
|
+
connections=(),
|
|
51
|
+
inputs=subgraph.inputs if self.is_wrapper() else (node.id,),
|
|
52
|
+
outputs=subgraph.outputs if self.is_wrapper() else (node.id,),
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
else:
|
|
56
|
+
node = self.to_node()
|
|
57
|
+
return Graph(nodes=(node,), connections=(), inputs=(node.id,), outputs=(node.id,))
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class Pipe[Input, Output, Metadata](BasePipe[Input, Output, Metadata]):
|
|
61
|
+
"""
|
|
62
|
+
Abstract base class of all synchronous pipes.
|
|
63
|
+
"""
|
|
64
|
+
|
|
65
|
+
def __or__[OtherOutput](self, other: "Pipe[Output, OtherOutput, Metadata]") -> "ChainPipe[Input, Output, OtherOutput, Metadata]":
|
|
66
|
+
"""
|
|
67
|
+
Alias for Pipe.chain.
|
|
68
|
+
"""
|
|
69
|
+
return self.chain(other)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
type PipeOrCallable[I, O, M] = Callable[[I, M], O] | BasePipe[I, O, M]
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def _call_or_apply[I, O, M](f: PipeOrCallable[I, O, M], d: I, m: M) -> O:
|
|
76
|
+
if isinstance(f, BasePipe):
|
|
77
|
+
return f.apply(d, m)
|
|
78
|
+
|
|
79
|
+
else:
|
|
80
|
+
return f(d, m)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
@dataclass(frozen=True)
|
|
84
|
+
class ChainPipe[Input, Middle, Output, Metadata](Pipe[Input, Output, Metadata]):
|
|
85
|
+
"""
|
|
86
|
+
Connects to pipes so that the second pipe is applied to the result of the first pipe.
|
|
87
|
+
"""
|
|
88
|
+
pipe1: BasePipe[Input, Middle, Metadata]
|
|
89
|
+
pipe2: BasePipe[Middle, Output, Metadata]
|
|
90
|
+
|
|
91
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
92
|
+
return self.pipe2.apply(self.pipe1.apply(data, metadata), metadata)
|
|
93
|
+
|
|
94
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
95
|
+
middle_batch = self.pipe1.batch_apply(data, metadata)
|
|
96
|
+
return self.pipe2.batch_apply(middle_batch, metadata)
|
|
97
|
+
|
|
98
|
+
def to_graph(self) -> Graph:
|
|
99
|
+
graph1 = self.pipe1.to_graph()
|
|
100
|
+
graph2 = self.pipe2.to_graph()
|
|
101
|
+
return graph1 >> graph2
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@dataclass(frozen=True)
|
|
105
|
+
class IdentityPipe[InputOutput, Metadata](Pipe[InputOutput, InputOutput, Metadata]):
|
|
106
|
+
"""
|
|
107
|
+
A pipe that returns its input.
|
|
108
|
+
"""
|
|
109
|
+
def apply(self, data: InputOutput, metadata: Metadata) -> InputOutput:
|
|
110
|
+
return data
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@dataclass(frozen=True)
|
|
114
|
+
class Merge2Pipe[Input, Output1, Output2, Output, Metadata](Pipe[Input, Output, Metadata]):
|
|
115
|
+
"""
|
|
116
|
+
Runs two pipes and merges their results.
|
|
117
|
+
"""
|
|
118
|
+
pipe1: BasePipe[Input, Output1, Metadata]
|
|
119
|
+
pipe2: BasePipe[Input, Output2, Metadata]
|
|
120
|
+
merge: Callable[[Output1, Output2], Output]
|
|
121
|
+
description: str = "Combine results"
|
|
122
|
+
|
|
123
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
124
|
+
output1 = self.pipe1.apply(data, metadata)
|
|
125
|
+
output2 = self.pipe2.apply(data, metadata)
|
|
126
|
+
return self.merge(output1, output2)
|
|
127
|
+
|
|
128
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
129
|
+
batch1 = self.pipe1.batch_apply(data, metadata)
|
|
130
|
+
batch2 = self.pipe2.batch_apply(data, metadata)
|
|
131
|
+
assert len(data) == len(batch1) == len(batch2), "lengths do not match"
|
|
132
|
+
return [self.merge(output1, output2) for output1, output2 in zip(batch1, batch2)]
|
|
133
|
+
|
|
134
|
+
def to_graph(self) -> Graph:
|
|
135
|
+
graph1 = self.pipe1.to_graph()
|
|
136
|
+
graph2 = self.pipe2.to_graph()
|
|
137
|
+
combine_node = self.to_node()._replace(title=self.description, shape="combine")
|
|
138
|
+
graph = Graph(nodes=(combine_node,), inputs=(combine_node.id,), outputs=(combine_node.id,))
|
|
139
|
+
return (graph1 | graph2) >> graph
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
@dataclass(frozen=True)
|
|
143
|
+
class Merge3Pipe[Input, Output1, Output2, Output3, Output, Metadata](Pipe[Input, Output, Metadata]):
|
|
144
|
+
"""
|
|
145
|
+
Runs three pipes and merges their results.
|
|
146
|
+
"""
|
|
147
|
+
pipe1: BasePipe[Input, Output1, Metadata]
|
|
148
|
+
pipe2: BasePipe[Input, Output2, Metadata]
|
|
149
|
+
pipe3: BasePipe[Input, Output3, Metadata]
|
|
150
|
+
merge: Callable[[Output1, Output2, Output3], Output]
|
|
151
|
+
description: str = "Combine results"
|
|
152
|
+
|
|
153
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
154
|
+
output1 = self.pipe1.apply(data, metadata)
|
|
155
|
+
output2 = self.pipe2.apply(data, metadata)
|
|
156
|
+
output3 = self.pipe3.apply(data, metadata)
|
|
157
|
+
return self.merge(output1, output2, output3)
|
|
158
|
+
|
|
159
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
160
|
+
batch1 = self.pipe1.batch_apply(data, metadata)
|
|
161
|
+
batch2 = self.pipe2.batch_apply(data, metadata)
|
|
162
|
+
batch3 = self.pipe3.batch_apply(data, metadata)
|
|
163
|
+
assert len(data) == len(batch1) == len(batch2) == len(batch3), "lengths do not match"
|
|
164
|
+
return [self.merge(output1, output2, output3) for output1, output2, output3 in zip(batch1, batch2, batch3)]
|
|
165
|
+
|
|
166
|
+
def to_graph(self) -> Graph:
|
|
167
|
+
graph1 = self.pipe1.to_graph()
|
|
168
|
+
graph2 = self.pipe2.to_graph()
|
|
169
|
+
graph3 = self.pipe3.to_graph()
|
|
170
|
+
combine_node = self.to_node()._replace(title=self.description, shape="combine")
|
|
171
|
+
graph = Graph(nodes=(combine_node,), inputs=(combine_node.id,), outputs=(combine_node.id,))
|
|
172
|
+
return (graph1 | graph2 | graph3) >> graph
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
@dataclass(frozen=True)
|
|
176
|
+
class BranchPipe[Input, Output, Metadata](Pipe[Input, Output, Metadata]):
|
|
177
|
+
"""
|
|
178
|
+
Selects one of two pipes to apply based on a condition.
|
|
179
|
+
"""
|
|
180
|
+
condition: PipeOrCallable[Input, bool, Metadata]
|
|
181
|
+
then_pipe: BasePipe[Input, Output, Metadata]
|
|
182
|
+
else_pipe: BasePipe[Input, Output, Metadata]
|
|
183
|
+
description: str = "Condition"
|
|
184
|
+
|
|
185
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
186
|
+
if _call_or_apply(self.condition, data, metadata):
|
|
187
|
+
return self.then_pipe.apply(data, metadata)
|
|
188
|
+
|
|
189
|
+
else:
|
|
190
|
+
return self.else_pipe.apply(data, metadata)
|
|
191
|
+
|
|
192
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
193
|
+
conds = [_call_or_apply(self.condition, d, metadata) for d in data]
|
|
194
|
+
then_batch = self.then_pipe.batch_apply([d for c, d in zip(conds, data) if c], metadata)
|
|
195
|
+
else_batch = self.else_pipe.batch_apply([d for c, d in zip(conds, data) if not c], metadata)
|
|
196
|
+
|
|
197
|
+
ans: list[Output] = []
|
|
198
|
+
t = 0
|
|
199
|
+
e = 0
|
|
200
|
+
for cond in conds:
|
|
201
|
+
if cond:
|
|
202
|
+
ans.append(then_batch[t])
|
|
203
|
+
t += 1
|
|
204
|
+
|
|
205
|
+
else:
|
|
206
|
+
ans.append(else_batch[e])
|
|
207
|
+
e += 1
|
|
208
|
+
|
|
209
|
+
return ans
|
|
210
|
+
|
|
211
|
+
def to_graph(self) -> Graph:
|
|
212
|
+
condition_node = self.to_node()._replace(title=self.description, shape="condition", subgraph=self.condition.to_graph() if isinstance(self.condition, Pipe) else None)
|
|
213
|
+
then_graph = self.then_pipe.to_graph()
|
|
214
|
+
else_graph = self.else_pipe.to_graph()
|
|
215
|
+
return (
|
|
216
|
+
(then_graph | else_graph)
|
|
217
|
+
.add(nodes=(condition_node,))
|
|
218
|
+
.add(connections=tuple(GraphConnection(condition_node.id, i, "yes") for i in then_graph.inputs))
|
|
219
|
+
.add(connections=tuple(GraphConnection(condition_node.id, i, "no") for i in else_graph.inputs))
|
|
220
|
+
._replace(
|
|
221
|
+
inputs=(condition_node.id,),
|
|
222
|
+
outputs=then_graph.outputs + else_graph.outputs,
|
|
223
|
+
)
|
|
224
|
+
)
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
@dataclass(frozen=True)
|
|
228
|
+
class ConditionalPipe[InputOutput, Metadata](Pipe[InputOutput, InputOutput, Metadata]):
|
|
229
|
+
"""
|
|
230
|
+
Applies the pipe only if the condition is met. Similar to BranchPipe except that the else branch is an identity function.
|
|
231
|
+
Can only be used when the subpipe returns the same type that gets in.
|
|
232
|
+
"""
|
|
233
|
+
condition: Callable[[InputOutput, Metadata], bool]
|
|
234
|
+
subpipe: BasePipe[InputOutput, InputOutput, Metadata]
|
|
235
|
+
description: str = "Condition"
|
|
236
|
+
|
|
237
|
+
def apply(self, data: InputOutput, metadata: Metadata) -> InputOutput:
|
|
238
|
+
if self.condition(data, metadata):
|
|
239
|
+
return self.subpipe.apply(data, metadata)
|
|
240
|
+
|
|
241
|
+
else:
|
|
242
|
+
return data
|
|
243
|
+
|
|
244
|
+
def batch_apply(self, data: Sequence[InputOutput], metadata: Metadata) -> Sequence[InputOutput]:
|
|
245
|
+
conds = [_call_or_apply(self.condition, d, metadata) for d in data]
|
|
246
|
+
then_batch = self.subpipe.batch_apply([d for c, d in zip(conds, data) if c], metadata)
|
|
247
|
+
|
|
248
|
+
ans: list[InputOutput] = []
|
|
249
|
+
t = 0
|
|
250
|
+
for i, cond in enumerate(conds):
|
|
251
|
+
if cond:
|
|
252
|
+
ans.append(then_batch[t])
|
|
253
|
+
t += 1
|
|
254
|
+
|
|
255
|
+
else:
|
|
256
|
+
ans.append(data[i])
|
|
257
|
+
|
|
258
|
+
return ans
|
|
259
|
+
|
|
260
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
261
|
+
return self.subpipe.to_graph()
|
|
262
|
+
|
|
263
|
+
def to_node(self) -> GraphNode:
|
|
264
|
+
return super().to_node()._replace(title=self.description)
|
|
265
|
+
|
|
266
|
+
def is_wrapper(self) -> bool:
|
|
267
|
+
return True
|
|
268
|
+
|
|
269
|
+
|
|
270
|
+
@dataclass(frozen=True)
|
|
271
|
+
class SelectPipe[Input, Output, Metadata, Key](Pipe[Input, Output, Metadata]):
|
|
272
|
+
"""
|
|
273
|
+
Selects one of multiple pipes based on a selector function. BranchPipe generalized to more than two branches.
|
|
274
|
+
"""
|
|
275
|
+
key: PipeOrCallable[Input, Key, Metadata]
|
|
276
|
+
subpipes: Mapping[Key, BasePipe[Input, Output, Metadata]]
|
|
277
|
+
otherwise_pipe: BasePipe[Input, Output, Metadata]
|
|
278
|
+
description: str = "Key"
|
|
279
|
+
|
|
280
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
281
|
+
key = _call_or_apply(self.key, data, metadata)
|
|
282
|
+
if key in self.subpipes:
|
|
283
|
+
return self.subpipes[key].apply(data, metadata)
|
|
284
|
+
|
|
285
|
+
else:
|
|
286
|
+
return self.otherwise_pipe.apply(data, metadata)
|
|
287
|
+
|
|
288
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
289
|
+
keys = [_call_or_apply(self.key, d, metadata) for d in data]
|
|
290
|
+
unique_keys = list(set(keys))
|
|
291
|
+
batches: dict[Key, list[tuple[int, Input]]] = {key: [] for key in unique_keys}
|
|
292
|
+
for i, (key, item) in enumerate(zip(keys, data)):
|
|
293
|
+
batches[key].append((i, item))
|
|
294
|
+
|
|
295
|
+
batch_results: list[Sequence[Output]] = []
|
|
296
|
+
for key in unique_keys:
|
|
297
|
+
pipe = self.subpipes.get(key, self.otherwise_pipe)
|
|
298
|
+
batch = [item for _, item in batches[key]]
|
|
299
|
+
batch_results.append(pipe.batch_apply(batch, metadata))
|
|
300
|
+
|
|
301
|
+
batch_results_dict = {key: batch_results[i] for i, key in enumerate(unique_keys)}
|
|
302
|
+
ans_dict = {i: result for key, batch in batches.items() for (i, _), result in zip(batch, batch_results_dict[key])}
|
|
303
|
+
return [ans_dict[i] for i in range(len(data))]
|
|
304
|
+
|
|
305
|
+
def to_graph(self) -> Graph:
|
|
306
|
+
condition_node = self.to_node()._replace(title=self.description, shape="condition", subgraph=self.key.to_graph() if isinstance(self.key, Pipe) else None)
|
|
307
|
+
otherwise_graph = self.otherwise_pipe.to_graph()
|
|
308
|
+
graph = otherwise_graph.add(
|
|
309
|
+
nodes=(condition_node,),
|
|
310
|
+
connections=tuple(GraphConnection(condition_node.id, i, "otherwise") for i in otherwise_graph.inputs),
|
|
311
|
+
)
|
|
312
|
+
for key, subpipe in self.subpipes.items():
|
|
313
|
+
subpipe_graph = subpipe.to_graph()
|
|
314
|
+
graph |= subpipe_graph
|
|
315
|
+
graph = graph.add(connections=tuple(GraphConnection(condition_node.id, i, str(key)) for i in subpipe_graph.inputs))
|
|
316
|
+
|
|
317
|
+
return graph._replace(inputs=(condition_node.id,))
|
|
318
|
+
|
|
319
|
+
|
|
320
|
+
@dataclass(frozen=True)
|
|
321
|
+
class ParallelPipe[Input, Output, CombinedOutput, Metadata](Pipe[Input, CombinedOutput, Metadata]):
|
|
322
|
+
"""
|
|
323
|
+
Executes multiple pipes on the same input and combines the outputs.
|
|
324
|
+
The synchronous implementation does not actually execute the pipes in parallel.
|
|
325
|
+
"""
|
|
326
|
+
subpipes: Sequence[BasePipe[Input, Output, Metadata]]
|
|
327
|
+
combine: PipeOrCallable[Sequence[Output], CombinedOutput, Metadata]
|
|
328
|
+
description: str = "Combine results"
|
|
329
|
+
|
|
330
|
+
def apply(self, data: Input, metadata: Metadata) -> CombinedOutput:
|
|
331
|
+
results: list[Output] = []
|
|
332
|
+
for subpipe in self.subpipes:
|
|
333
|
+
results.append(subpipe.apply(data, metadata))
|
|
334
|
+
|
|
335
|
+
seq: Sequence[Output] = results
|
|
336
|
+
return _call_or_apply(self.combine, seq, metadata)
|
|
337
|
+
|
|
338
|
+
def batch_apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[CombinedOutput]:
|
|
339
|
+
results: list[Sequence[Output]] = []
|
|
340
|
+
for subpipe in self.subpipes:
|
|
341
|
+
results.append(subpipe.batch_apply(data, metadata))
|
|
342
|
+
|
|
343
|
+
zipped_results: zip[Sequence[Output]] = zip(*results)
|
|
344
|
+
return [_call_or_apply(self.combine, seq, metadata) for seq in zipped_results]
|
|
345
|
+
|
|
346
|
+
def to_graph(self) -> Graph:
|
|
347
|
+
combine_node = self.to_node()._replace(title=self.description, shape="combine", subgraph=self.combine.to_graph() if isinstance(self.combine, Pipe) else None)
|
|
348
|
+
graph = Graph(nodes=(combine_node,))
|
|
349
|
+
for subpipe in self.subpipes:
|
|
350
|
+
subpipe_graph = subpipe.to_graph()
|
|
351
|
+
graph |= subpipe_graph
|
|
352
|
+
graph = graph.add(connections=tuple(GraphConnection(o, combine_node.id) for o in subpipe_graph.outputs))
|
|
353
|
+
|
|
354
|
+
return graph._replace(outputs=(combine_node.id,))
|
|
355
|
+
|
|
356
|
+
|
|
357
|
+
@dataclass(frozen=True)
|
|
358
|
+
class MetadataWrapperPipe[Input, Output, OuterMetadata, InnerMetadata](Pipe[Input, Output, OuterMetadata]):
|
|
359
|
+
"""
|
|
360
|
+
Executes a pipe with changed metadata.
|
|
361
|
+
|
|
362
|
+
Does NOT allow to_inner to be a pipe: this is intentional, metadata is not transformed by pipes.
|
|
363
|
+
"""
|
|
364
|
+
to_inner: Callable[[OuterMetadata], InnerMetadata]
|
|
365
|
+
subpipe: BasePipe[Input, Output, InnerMetadata]
|
|
366
|
+
description: str = "Wrap metadata"
|
|
367
|
+
|
|
368
|
+
def apply(self, data: Input, metadata: OuterMetadata) -> Output:
|
|
369
|
+
return self.subpipe.apply(data, self.to_inner(metadata))
|
|
370
|
+
|
|
371
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
372
|
+
return self.subpipe.to_graph()
|
|
373
|
+
|
|
374
|
+
def to_node(self) -> GraphNode:
|
|
375
|
+
return super().to_node()._replace(title=self.description)
|
|
376
|
+
|
|
377
|
+
def is_wrapper(self) -> bool:
|
|
378
|
+
return True
|
|
379
|
+
|
|
380
|
+
|
|
381
|
+
@dataclass(frozen=True)
|
|
382
|
+
class MapPipe[Input, Output, Metadata](Pipe[Sequence[Input], Sequence[Output], Metadata]):
|
|
383
|
+
"""
|
|
384
|
+
Maps a sequence of input objects into a sequence of output objects with a subpipe.
|
|
385
|
+
|
|
386
|
+
Does NOT allow the subpipe to be a callable: this is intentional, functions that transforms Inputs to Outputs should be pipes.
|
|
387
|
+
"""
|
|
388
|
+
subpipe: BasePipe[Input, Output, Metadata]
|
|
389
|
+
description: str = "Map"
|
|
390
|
+
|
|
391
|
+
def apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Output]:
|
|
392
|
+
return self.subpipe.batch_apply(data, metadata)
|
|
393
|
+
|
|
394
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
395
|
+
return self.subpipe.to_graph()
|
|
396
|
+
|
|
397
|
+
def to_node(self) -> GraphNode:
|
|
398
|
+
return super().to_node()._replace(title=self.description)
|
|
399
|
+
|
|
400
|
+
|
|
401
|
+
@dataclass(frozen=True)
|
|
402
|
+
class FilterPipe[Input, Metadata](Pipe[Sequence[Input], Sequence[Input], Metadata]):
|
|
403
|
+
"""
|
|
404
|
+
Filters a sequence of input object with a predicate.
|
|
405
|
+
"""
|
|
406
|
+
predicate: PipeOrCallable[Input, bool, Metadata]
|
|
407
|
+
description: str = "Filter"
|
|
408
|
+
|
|
409
|
+
def apply(self, data: Sequence[Input], metadata: Metadata) -> Sequence[Input]:
|
|
410
|
+
return [i for i in data if _call_or_apply(self.predicate, i, metadata)]
|
|
411
|
+
|
|
412
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
413
|
+
return self.predicate.to_graph() if isinstance(self.predicate, Pipe) else None
|
|
414
|
+
|
|
415
|
+
def to_node(self) -> GraphNode:
|
|
416
|
+
return super().to_node()._replace(title=self.description)
|
|
417
|
+
|
|
418
|
+
|
|
419
|
+
@dataclass(frozen=True)
|
|
420
|
+
class RetryPipe[Input, Output, Metadata](Pipe[Input, Output, Metadata]):
|
|
421
|
+
"""
|
|
422
|
+
Tries to execute the subpipe several times in case it fails.
|
|
423
|
+
"""
|
|
424
|
+
subpipe: BasePipe[Input, Output, Metadata]
|
|
425
|
+
retries: int
|
|
426
|
+
exceptions: type | tuple[type, ...]
|
|
427
|
+
description: str = "Retry several times"
|
|
428
|
+
|
|
429
|
+
def apply(self, data: Input, metadata: Metadata) -> Output:
|
|
430
|
+
attempt = 0
|
|
431
|
+
while True:
|
|
432
|
+
try:
|
|
433
|
+
return self.subpipe.apply(data, metadata)
|
|
434
|
+
|
|
435
|
+
except Exception as e:
|
|
436
|
+
if not isinstance(e, self.exceptions):
|
|
437
|
+
raise e
|
|
438
|
+
|
|
439
|
+
attempt += 1
|
|
440
|
+
if attempt > self.retries:
|
|
441
|
+
raise e
|
|
442
|
+
|
|
443
|
+
def get_subgraph(self) -> Optional[Graph]:
|
|
444
|
+
return self.subpipe.to_graph()
|
|
445
|
+
|
|
446
|
+
def to_node(self) -> GraphNode:
|
|
447
|
+
return super().to_node()._replace(title=self.description)
|
|
448
|
+
|
|
449
|
+
def is_wrapper(self) -> bool:
|
|
450
|
+
return True
|
|
451
|
+
|
|
File without changes
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: kpipeline
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: A pipeline library for Python
|
|
5
|
+
Requires-Python: >=3.13
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
|
|
8
|
+
# kpipe
|
|
9
|
+
|
|
10
|
+
`kpipe` is a simple pipeline library that allows you to write your complex application as a pipeline. This helps code organization, debugging, and testing.
|
|
11
|
+
|
|
12
|
+
## How to use
|
|
13
|
+
|
|
14
|
+
kpipe is like a DSL language inside Python. It is based on "pipes" which can be combined into pipelines using the pipe primitives provided by the library.
|
|
15
|
+
|
|
16
|
+
Each pipe represents a function transforming an input into an output. Pipes are stateless and immutable. Below is an example defining two example pipes.
|
|
17
|
+
|
|
18
|
+
```python
|
|
19
|
+
from kpipeline import Pipe
|
|
20
|
+
|
|
21
|
+
class AddOnePipe(Pipe[int, int, None]):
|
|
22
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
23
|
+
return input + 1
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class MulByTwoPipe(Pipe[int, int, None]):
|
|
27
|
+
def apply(self, input: int, metadata: None) -> int:
|
|
28
|
+
return input * 2
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
These pipes can be combined using the ChainPipe primitive (aliased into the `|` operator) two form a pipeline that performs these two pipes sequentially:
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
pipeline = AddOnePipe() | MulByTwoPipe()
|
|
35
|
+
|
|
36
|
+
# or
|
|
37
|
+
|
|
38
|
+
from kpipeline import ChainPipe
|
|
39
|
+
pipeline = ChainPipe(AddOnePipe(), MulByTwoPipe())
|
|
40
|
+
|
|
41
|
+
print(pipeline.apply(2)) # 6
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
Currently, the library defines these primitives:
|
|
45
|
+
|
|
46
|
+
| Primitive | Purpose |
|
|
47
|
+
|---|---|
|
|
48
|
+
|`ChainPipe`|Execute two pipes sequentially|
|
|
49
|
+
|`ConditionalPipe`|Execute a pipe if a given condition is true|
|
|
50
|
+
|`BranchPipe`|Execute one of two pipes depending on whether the condition is true or false|
|
|
51
|
+
|`SelectPipe`|Execute one of multiple pipes based on a selector key|
|
|
52
|
+
|`ParallelPipe`|Execute multiple pipes and combine their results (only async implementation is actually parallel)|
|
|
53
|
+
|`MetadataWrapperPipe`|Transform the metadata given to the pipeline for the subpipe|
|
|
54
|
+
|`MapPipe`|Apply a subpipe into a sequence of inputs|
|
|
55
|
+
|`FilterPipe`|Filter a sequence of inputs using a predicate|
|
|
56
|
+
|`RetryPipe`|Run a pipe multiple times in case it fails|
|
|
57
|
+
|
|
58
|
+
## Tests
|
|
59
|
+
|
|
60
|
+
You can run all tests in this repository with
|
|
61
|
+
|
|
62
|
+
```
|
|
63
|
+
uv run pytest test
|
|
64
|
+
uv run mypy .
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## License
|
|
68
|
+
|
|
69
|
+
TBD
|
|
70
|
+
|
|
71
|
+
## AI Use Disclosure
|
|
72
|
+
|
|
73
|
+
I hate writing tests so unit tests have been generated by GPT-OSS 120B for your convenience.
|
|
74
|
+
They are provided in the hope that they are better than nothing.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
kpipeline/__init__.py
|
|
4
|
+
kpipeline/async_pipeline.py
|
|
5
|
+
kpipeline/graph.py
|
|
6
|
+
kpipeline/pipeline.py
|
|
7
|
+
kpipeline/py.typed
|
|
8
|
+
kpipeline.egg-info/PKG-INFO
|
|
9
|
+
kpipeline.egg-info/SOURCES.txt
|
|
10
|
+
kpipeline.egg-info/dependency_links.txt
|
|
11
|
+
kpipeline.egg-info/top_level.txt
|
|
12
|
+
test/test_graph.py
|
|
13
|
+
test/test_graph_dot.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
kpipeline
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "kpipeline"
|
|
3
|
+
version = "0.1.1"
|
|
4
|
+
description = "A pipeline library for Python"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
dependencies = []
|
|
8
|
+
|
|
9
|
+
[dependency-groups]
|
|
10
|
+
dev = [
|
|
11
|
+
"mypy>=1.18.2",
|
|
12
|
+
"pytest>=8.4.2",
|
|
13
|
+
"pytest-asyncio>=1.2.0",
|
|
14
|
+
]
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
import builtins
|
|
2
|
+
import types
|
|
3
|
+
import pytest
|
|
4
|
+
|
|
5
|
+
from kpipeline.graph import (
|
|
6
|
+
Graph,
|
|
7
|
+
GraphNode,
|
|
8
|
+
GraphConnection,
|
|
9
|
+
)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
# ----------------------------------------------------------------------
|
|
13
|
+
# Graph unit‑tests
|
|
14
|
+
# ----------------------------------------------------------------------
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def test_graph_addition_and_validity():
|
|
18
|
+
"""`Graph.__or__` must concatenate the underlying tuples."""
|
|
19
|
+
n1 = GraphNode(id="a", title="A")
|
|
20
|
+
n2 = GraphNode(id="b", title="B")
|
|
21
|
+
c1 = GraphConnection("a", "b", "link")
|
|
22
|
+
|
|
23
|
+
g1 = Graph(nodes=(n1,), connections=(c1,), inputs=("a",), outputs=("b",))
|
|
24
|
+
g2 = Graph(nodes=(n2,), connections=(), inputs=(), outputs=())
|
|
25
|
+
|
|
26
|
+
g3 = g1 | g2
|
|
27
|
+
|
|
28
|
+
assert isinstance(g3, Graph)
|
|
29
|
+
assert g3.nodes == (n1, n2)
|
|
30
|
+
assert g3.connections == (c1,)
|
|
31
|
+
assert g3.inputs == ("a",)
|
|
32
|
+
assert g3.outputs == ("b",)
|
|
33
|
+
|
|
34
|
+
# The graph is valid because every reference points to an existing node.
|
|
35
|
+
assert g3.is_valid()
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def test_graph_rshift_creates_cross_connections():
|
|
39
|
+
"""`Graph.__rshift__` must connect every output of the left graph to every input of the right graph."""
|
|
40
|
+
n_left = GraphNode(id="l", title="L")
|
|
41
|
+
n_right = GraphNode(id="r", title="R")
|
|
42
|
+
|
|
43
|
+
left = Graph(nodes=(n_left,), connections=(), inputs=("l",), outputs=("l",))
|
|
44
|
+
right = Graph(nodes=(n_right,), connections=(), inputs=("r",), outputs=("r",))
|
|
45
|
+
|
|
46
|
+
composed = left >> right
|
|
47
|
+
|
|
48
|
+
# Two original nodes plus the automatically created connection
|
|
49
|
+
assert composed.nodes == (n_left, n_right)
|
|
50
|
+
assert len(composed.connections) == 1
|
|
51
|
+
conn = composed.connections[0]
|
|
52
|
+
assert conn.from_node == "l"
|
|
53
|
+
assert conn.to_node == "r"
|
|
54
|
+
assert conn.label == "" # default label
|
|
55
|
+
|
|
56
|
+
# Inputs come from the left side, outputs from the right side
|
|
57
|
+
assert composed.inputs == ("l",)
|
|
58
|
+
assert composed.outputs == ("r",)
|
|
59
|
+
|
|
60
|
+
# The composed graph is still valid
|
|
61
|
+
assert composed.is_valid()
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_graph_is_invalid_when_reference_missing():
|
|
65
|
+
n = GraphNode(id="x", title="X")
|
|
66
|
+
g = Graph(nodes=(n,), connections=(GraphConnection("x", "y"),), inputs=("x",), outputs=("x",))
|
|
67
|
+
assert not g.is_valid() # "y" does not exist
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
from kpipeline.graph import Graph, GraphNode, GraphConnection
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_simple_graph_dot():
|
|
5
|
+
node1 = GraphNode(id="n1", title="Node 1")
|
|
6
|
+
node2 = GraphNode(id="n2", title="Node 2")
|
|
7
|
+
conn = GraphConnection(from_node="n1", to_node="n2", label="connect")
|
|
8
|
+
graph = Graph(nodes=(node1, node2), connections=(conn,))
|
|
9
|
+
|
|
10
|
+
dot = graph.to_dot()
|
|
11
|
+
assert 'digraph graph {' in dot
|
|
12
|
+
assert '"n1" [label="Node 1", shape=box];' in dot
|
|
13
|
+
assert '"n2" [label="Node 2", shape=box];' in dot
|
|
14
|
+
assert '"n1" -> "n2" [label="connect"];' in dot
|
|
15
|
+
|
|
16
|
+
def test_shapes_dot():
|
|
17
|
+
node_combine = GraphNode(id="nc", title="Combine", shape="combine")
|
|
18
|
+
node_cond = GraphNode(id="nd", title="Condition", shape="condition")
|
|
19
|
+
graph = Graph(nodes=(node_combine, node_cond))
|
|
20
|
+
|
|
21
|
+
dot = graph.to_dot()
|
|
22
|
+
assert '"nc" [label="Combine", shape=ellipse];' in dot
|
|
23
|
+
assert '"nd" [label="Condition", shape=diamond];' in dot
|
|
24
|
+
|
|
25
|
+
def test_subgraph_dot():
|
|
26
|
+
inner_node = GraphNode(id="in", title="Inner")
|
|
27
|
+
inner_graph = Graph(nodes=(inner_node,))
|
|
28
|
+
|
|
29
|
+
outer_node = GraphNode(id="out", title="Outer", subgraph=inner_graph)
|
|
30
|
+
graph = Graph(nodes=(outer_node,))
|
|
31
|
+
|
|
32
|
+
dot = graph.to_dot()
|
|
33
|
+
assert 'subgraph "cluster_out" {' in dot
|
|
34
|
+
assert 'label="Outer";' in dot
|
|
35
|
+
assert '"in" [label="Inner", shape=box];' in dot
|