118 lines
2.8 KiB
Python
118 lines
2.8 KiB
Python
from dataclasses import dataclass
|
|
|
|
from src.breakshaft.convertor import ConvRepo
|
|
|
|
|
|
@dataclass
|
|
class A:
|
|
a: int
|
|
|
|
|
|
@dataclass
|
|
class B:
|
|
b: float
|
|
|
|
|
|
type optC = str
|
|
|
|
|
|
def test_default_consumer_args():
|
|
repo = ConvRepo(store_sources=True)
|
|
|
|
@repo.mark_injector()
|
|
def b_to_a(b: B) -> A:
|
|
return A(int(b.b))
|
|
|
|
@repo.mark_injector()
|
|
def a_to_b(a: A) -> B:
|
|
return B(float(a.a))
|
|
|
|
@repo.mark_injector()
|
|
def int_to_a(i: int) -> A:
|
|
return A(i)
|
|
|
|
type ret1 = tuple[int, str]
|
|
|
|
def consumer1(dep: A, opt_dep: optC = '42') -> ret1:
|
|
return dep.a, opt_dep
|
|
|
|
def consumer2(dep: A, dep1: ret1) -> optC:
|
|
return str((dep.a, dep1))
|
|
|
|
p1 = repo.create_pipeline(
|
|
(B,),
|
|
[consumer1, consumer2],
|
|
force_commutative=True,
|
|
allow_sync=True,
|
|
allow_async=False,
|
|
force_async=False
|
|
)
|
|
res = p1(B(42.1))
|
|
assert res == "(42, (42, '42'))"
|
|
|
|
p2 = repo.create_pipeline(
|
|
(B,),
|
|
[consumer1, consumer2, consumer1],
|
|
force_commutative=True,
|
|
allow_sync=True,
|
|
allow_async=False,
|
|
force_async=False
|
|
)
|
|
res = p2(B(42.1))
|
|
assert res == (42, "(42, (42, '42'))")
|
|
|
|
|
|
def test_pipeline_with_subgraph_duplicates():
|
|
repo = ConvRepo()
|
|
|
|
b_to_a_calls = [0]
|
|
|
|
@repo.mark_injector()
|
|
def b_to_a(b: B) -> A:
|
|
b_to_a_calls[0] += 1
|
|
return A(int(b.b))
|
|
|
|
@repo.mark_injector()
|
|
def a_to_b(a: A) -> B:
|
|
return B(float(a.a))
|
|
|
|
@repo.mark_injector()
|
|
def int_to_a(i: int) -> A:
|
|
return A(i)
|
|
|
|
type ret1 = tuple[int, str]
|
|
|
|
cons1_calls = [0]
|
|
cons2_calls = [0]
|
|
|
|
def consumer1(dep: A, opt_dep: optC = '42') -> A:
|
|
cons1_calls[0] += 1
|
|
return A(dep.a + int(opt_dep))
|
|
|
|
def consumer2(dep: A) -> optC:
|
|
cons2_calls[0] += 1
|
|
return str(dep.a)
|
|
|
|
p1 = repo.create_pipeline(
|
|
(B,),
|
|
[consumer1, consumer2, consumer1, consumer2, consumer1, consumer2, consumer1, consumer2, consumer1],
|
|
force_commutative=True,
|
|
allow_sync=True,
|
|
allow_async=False,
|
|
force_async=False
|
|
)
|
|
res = p1(B(42.1))
|
|
assert res.a == 42 + (42 * 31)
|
|
assert b_to_a_calls[0] == 1
|
|
assert cons1_calls[0] == 5
|
|
assert cons2_calls[0] == 4
|
|
|
|
|
|
def convertor(_5891515089754: "<class 'test_pipeline.B'>"):
|
|
# <function test_default_consumer_args.<locals>.b_to_a at 0x7f5bb1be02c0>
|
|
_5891515089643 = _conv_funcmap[8751987548204](b=_5891515089754)
|
|
# <function test_default_consumer_args.<locals>.consumer1 at 0x7f5bb1be0c20>
|
|
_8751987542640 = _conv_funcmap[8751987548354](dep=_5891515089643)
|
|
# <function test_default_consumer_args.<locals>.consumer2 at 0x7f5bb1be0540>
|
|
_8751987537115 = _conv_funcmap[8751987548244](dep=_5891515089643, dep1=_8751987542640)
|
|
return _8751987542640 |