feat: add StagedPipelineBuilder with stage-based composition · dexpace/python-sdk@93e02c0 · GitHub
Skip to content

Commit 93e02c0

Browse files
committed
feat: add StagedPipelineBuilder with stage-based composition
1 parent 794b704 commit 93e02c0

5 files changed

Lines changed: 680 additions & 0 deletions

File tree

‎packages/dexpace-sdk-core/src/dexpace/sdk/core/pipeline/__init__.py‎

Lines changed: 6 additions & 0 deletions
Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
"""Async twin of :class:`StagedPipelineBuilder`.
2+
3+
Behaviour mirrors the sync builder exactly: pillar enforcement, surgical
4+
edits, ``from_pipeline`` round-trip. The only differences are the policy
5+
types (``AsyncPolicy`` instead of ``Policy``) and the produced pipeline
6+
(``AsyncPipeline`` instead of ``Pipeline``).
7+
"""
8+
9+
from __future__ import annotations
10+
11+
from typing import TYPE_CHECKING, Self
12+
13+
from .async_pipeline import AsyncPipeline
14+
from .async_policy import AsyncPolicy
15+
from .stage import Stage
16+
17+
if TYPE_CHECKING:
18+
from ..client.async_http_client import AsyncHttpClient
19+
20+
21+
class AsyncStagedPipelineBuilder:
22+
"""Build an :class:`AsyncPipeline` by stage rather than user-specified order.
23+
24+
See :class:`StagedPipelineBuilder` for the full behaviour; this is the
25+
async counterpart.
26+
"""
27+
28+
__slots__ = ("_buckets", "_client", "_pillars")
29+
30+
def __init__(self, client: AsyncHttpClient) -> None:
31+
self._client = client
32+
self._pillars: dict[Stage, AsyncPolicy] = {}
33+
self._buckets: dict[Stage, list[AsyncPolicy]] = {}
34+
35+
def append(self, policy: AsyncPolicy, *, force: bool = False) -> Self:
36+
"""Append ``policy`` to the tail of its stage's bucket."""
37+
stage = policy.STAGE
38+
if stage.is_pillar:
39+
self._install_pillar(policy, stage, force=force)
40+
else:
41+
self._buckets.setdefault(stage, []).append(policy)
42+
return self
43+
44+
def prepend(self, policy: AsyncPolicy, *, force: bool = False) -> Self:
45+
"""Prepend ``policy`` to the head of its stage's bucket."""
46+
stage = policy.STAGE
47+
if stage.is_pillar:
48+
self._install_pillar(policy, stage, force=force)
49+
else:
50+
self._buckets.setdefault(stage, []).insert(0, policy)
51+
return self
52+
53+
def replace(self, target: type[AsyncPolicy], new: AsyncPolicy) -> Self:
54+
"""Replace the first instance of ``target`` with ``new``."""
55+
for stage, pillar in self._pillars.items():
56+
if isinstance(pillar, target):
57+
del self._pillars[stage]
58+
self.append(new, force=True)
59+
return self
60+
for stage, bucket in self._buckets.items():
61+
for i, p in enumerate(bucket):
62+
if isinstance(p, target):
63+
if stage == new.STAGE:
64+
bucket[i] = new
65+
else:
66+
del bucket[i]
67+
self.append(new)
68+
return self
69+
raise ValueError(f"No instance of {target.__name__} in the builder.")
70+
71+
def insert_after(self, target: type[AsyncPolicy], new: AsyncPolicy) -> Self:
72+
"""Insert ``new`` immediately after the first ``target`` instance."""
73+
return self._splice(target, new, offset=1)
74+
75+
def insert_before(self, target: type[AsyncPolicy], new: AsyncPolicy) -> Self:
76+
"""Insert ``new`` immediately before the first ``target`` instance."""
77+
return self._splice(target, new, offset=0)
78+
79+
def remove(self, target: type[AsyncPolicy]) -> Self:
80+
"""Remove every instance of ``target`` from the builder."""
81+
self._pillars = {s: p for s, p in self._pillars.items() if not isinstance(p, target)}
82+
for stage in list(self._buckets):
83+
self._buckets[stage] = [p for p in self._buckets[stage] if not isinstance(p, target)]
84+
if not self._buckets[stage]:
85+
del self._buckets[stage]
86+
return self
87+
88+
def build(self) -> AsyncPipeline:
89+
"""Flatten the builder's contents into an :class:`AsyncPipeline`."""
90+
return AsyncPipeline(self._client, policies=self._flatten())
91+
92+
@classmethod
93+
def from_pipeline(cls, pipeline: AsyncPipeline) -> Self:
94+
"""Seed a builder from an existing :class:`AsyncPipeline`."""
95+
from ._async_transport_runner import _AsyncTransportRunner
96+
97+
builder = cls(pipeline.transport)
98+
chain: list[AsyncPolicy] = []
99+
node: AsyncPolicy | None = pipeline._chain
100+
while node is not None and not isinstance(node, _AsyncTransportRunner):
101+
chain.append(node)
102+
node = getattr(node, "next", None)
103+
last_stage: Stage | None = None
104+
for policy in chain:
105+
if last_stage is not None and last_stage > policy.STAGE:
106+
raise ValueError(
107+
f"AsyncPipeline policy {type(policy).__name__} at stage {policy.STAGE} "
108+
f"comes after stage {last_stage}; staged builder requires "
109+
f"non-decreasing stage order. Use the list constructor instead."
110+
)
111+
last_stage = policy.STAGE
112+
builder.append(policy, force=True)
113+
return builder
114+
115+
def _install_pillar(self, policy: AsyncPolicy, stage: Stage, *, force: bool) -> None:
116+
if stage in self._pillars and not force:
117+
existing = type(self._pillars[stage]).__name__
118+
raise ValueError(
119+
f"Pillar stage {stage.name} is already filled by {existing}. "
120+
f"Use replace({type(policy).__name__}, new) to swap, or "
121+
f"force=True to overwrite."
122+
)
123+
self._pillars[stage] = policy
124+
125+
def _splice(self, target: type[AsyncPolicy], new: AsyncPolicy, *, offset: int) -> Self:
126+
flat = self._flatten()
127+
for i, p in enumerate(flat):
128+
if isinstance(p, target):
129+
flat.insert(i + offset, new)
130+
self._reload(flat)
131+
return self
132+
raise ValueError(f"No instance of {target.__name__} in the builder.")
133+
134+
def _flatten(self) -> list[AsyncPolicy]:
135+
out: list[AsyncPolicy] = []
136+
for stage in Stage:
137+
if stage is Stage.SEND:
138+
continue
139+
if stage.is_pillar:
140+
if stage in self._pillars:
141+
out.append(self._pillars[stage])
142+
else:
143+
out.extend(self._buckets.get(stage, ()))
144+
return out
145+
146+
def _reload(self, policies: list[AsyncPolicy]) -> None:
147+
self._pillars.clear()
148+
self._buckets.clear()
149+
for p in policies:
150+
if p.STAGE.is_pillar:
151+
self._pillars[p.STAGE] = p
152+
else:
153+
self._buckets.setdefault(p.STAGE, []).append(p)
154+
155+
156+
__all__ = ["AsyncStagedPipelineBuilder"]
Lines changed: 212 additions & 0 deletions

0 commit comments

Comments
 (0)