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.
@@ -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
+ 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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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