Skip to content

Commit b4c3cdd

Browse files
tzaffitzaffi
authored
Refactor graviton usage (#628)
1 parent 4c716b4 commit b4c3cdd

File tree

8 files changed

+104
-154
lines changed

8 files changed

+104
-154
lines changed

‎requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
black==22.3.0
22
flake8==5.0.4
33
flake8-tidy-imports==4.6.0
4-
graviton@git+https://github.com/algorand/graviton@v0.7.1
4+
graviton@git+https://github.com/algorand/graviton@v0.8.0
55
mypy==0.991
66
pytest==7.2.0
77
pytest-cov==3.0.0

‎tests/blackbox.py

Lines changed: 49 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
1-
from typing import Any, Callable, Generic, Optional, Sequence, TypeVar, cast
2-
from dataclasses import dataclass
1+
from typing import Any, Callable, Sequence, cast
32

43
import algosdk.abi
54
from algosdk.v2client import algod
65

76
from graviton import blackbox
8-
from graviton.blackbox import DryRunInspector, DryRunExecutor
7+
from graviton.blackbox import (
8+
DryRunInspector,
9+
DryRunExecutor,
10+
DryRunTransactionParams as TxParams,
11+
)
912
from graviton.models import PyTypes
1013

1114
from pyteal.ast.subroutine import OutputKwArgInfo
@@ -138,37 +141,13 @@ def decorator_blackbox(func: SubroutineFnWrapper | ABIReturnSubroutine):
138141
# ---- API ---- #
139142

140143

141-
Output = TypeVar("Output")
142-
Lazy = Callable[[], Output]
143-
144-
145-
@dataclass(frozen=True)
146-
class _MatchMode(Generic[Output]):
147-
runner: Optional["PyTealDryRunExecutor"]
148-
app_case: Lazy
149-
signature_case: Lazy
150-
trace: Any = None
151-
152-
def __post_init__(self):
153-
if self.runner and self.trace:
154-
self.runner.add_trace(self.trace)
155-
156-
def __call__(self, mode: Mode, *args, **kwargs) -> Output:
157-
match mode:
158-
case Mode.Application:
159-
return self.app_case()
160-
case Mode.Signature:
161-
return self.signature_case()
162-
case _:
163-
raise Exception(f"Unknown mode {mode} of type {type(mode)}")
164-
165-
166144
def mode_to_execution_mode(mode: Mode) -> blackbox.ExecutionMode:
167-
return _MatchMode(
168-
runner=None,
169-
app_case=lambda: blackbox.ExecutionMode.Application,
170-
signature_case=lambda: blackbox.ExecutionMode.Signature,
171-
)(mode)
145+
if mode == Mode.Application:
146+
return blackbox.ExecutionMode.Application
147+
if mode == Mode.Signature:
148+
return blackbox.ExecutionMode.Signature
149+
150+
raise ValueError(f"Can't handle {mode=}")
172151

173152

174153
class PyTealDryRunExecutor:
@@ -182,9 +161,11 @@ def __init__(self, subr: BlackboxWrapper, mode: Mode):
182161
mode: type of program to produce: logic sig (Mode.Signature) or app (Mode.Application)
183162
"""
184163
input_types = subr.input_types
185-
assert (
186-
input_types is not None
187-
), "please provide input_types in your @Subroutine or @ABIReturnSubroutine annotation (this is crucial for generating proper end-to-end testable PyTeal)"
164+
assert input_types is not None, (
165+
"please provide input_types in your @Subroutine or @ABIReturnSubroutine "
166+
"annotation. "
167+
"(this is crucial for generating proper end-to-end testable PyTeal)"
168+
)
188169

189170
self.subr, self.mode, self.input_types = subr, mode, input_types
190171
match subr.subroutine:
@@ -417,68 +398,45 @@ def approval():
417398
return approval
418399

419400
def compile(self, version: int, assemble_constants: bool = False) -> str:
420-
return _MatchMode(
421-
runner=self,
422-
app_case=lambda: compileTeal(
423-
self.program(),
424-
self.mode,
425-
version=version,
426-
assembleConstants=assemble_constants,
427-
),
428-
signature_case=lambda: compileTeal(
429-
self.program(),
430-
self.mode,
431-
version=version,
432-
assembleConstants=assemble_constants,
433-
),
434-
)(self.mode)
435-
436-
def dryrun_on_sequence(
401+
return compileTeal(
402+
self.program(),
403+
self.mode,
404+
version=version,
405+
assembleConstants=assemble_constants,
406+
)
407+
408+
def executor(self, compiler_version: int = 6) -> DryRunExecutor:
409+
return DryRunExecutor(
410+
algod=algod_with_assertion(),
411+
mode=mode_to_execution_mode(self.mode),
412+
teal=self.compile(compiler_version),
413+
abi_method_signature=self.abi_method_signature(),
414+
omit_method_selector=True,
415+
)
416+
417+
def dryrun_sequence(
437418
self,
438419
inputs: list[Sequence[PyTypes]],
420+
*,
439421
compiler_version=6,
422+
txn_params: TxParams | None = None,
423+
verbose: bool = False,
440424
) -> list[DryRunInspector]:
441-
teal = self.compile(compiler_version)
442-
return _MatchMode(
443-
self,
444-
app_case=lambda: DryRunExecutor.dryrun_app_on_sequence(
445-
algod=algod_with_assertion(),
446-
teal=teal,
447-
inputs=inputs,
448-
abi_method_signature=self.abi_method_signature(),
449-
omit_method_selector=True,
425+
return cast(
426+
list,
427+
self.executor(compiler_version).run_sequence(
428+
inputs, txn_params=txn_params, verbose=verbose
450429
),
451-
signature_case=lambda: DryRunExecutor.dryrun_logicsig_on_sequence(
452-
algod=algod_with_assertion(),
453-
teal=teal,
454-
inputs=inputs,
455-
abi_method_signature=self.abi_method_signature(),
456-
omit_method_selector=True,
457-
),
458-
trace=teal,
459-
)(self.mode)
430+
)
460431

461-
def dryrun(
432+
def dryrun_one(
462433
self,
463434
args: Sequence[bytes | str | int],
435+
*,
464436
compiler_version=6,
437+
txn_params: TxParams | None = None,
438+
verbose: bool = False,
465439
) -> DryRunInspector:
466-
teal = self.compile(compiler_version)
467-
return _MatchMode(
468-
self,
469-
app_case=lambda: DryRunExecutor.dryrun_app(
470-
algod_with_assertion(),
471-
teal,
472-
args,
473-
abi_method_signature=self.abi_method_signature(),
474-
omit_method_selector=True,
475-
),
476-
signature_case=lambda: DryRunExecutor.dryrun_logicsig(
477-
algod_with_assertion(),
478-
teal,
479-
args,
480-
abi_method_signature=self.abi_method_signature(),
481-
omit_method_selector=True,
482-
),
483-
trace=teal,
484-
)(self.mode)
440+
return self.executor(compiler_version).run_one(
441+
args, txn_params=txn_params, verbose=verbose
442+
)

‎tests/integration/abi_roundtrip_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ def test_roundtrip(abi_type):
257257
args = (rand_abi_instance,)
258258

259259
def dryrun_roundtrip(version: int):
260-
inspector = roundtripper.dryrun(args, compiler_version=version)
260+
inspector = roundtripper.dryrun_one(args, compiler_version=version)
261261

262262
cost = inspector.cost()
263263
passed = inspector.passed()

‎tests/integration/ecdsa_test.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def verify():
4646
)
4747

4848
args = []
49-
app_result = PyTealDryRunExecutor(verify, Mode.Application).dryrun(
49+
app_result = PyTealDryRunExecutor(verify, Mode.Application).dryrun_one(
5050
args, compiler_version=5
5151
)
5252

@@ -81,7 +81,7 @@ def verify_fail():
8181
)
8282

8383
args = []
84-
app_result = PyTealDryRunExecutor(verify_fail, Mode.Application).dryrun(
84+
app_result = PyTealDryRunExecutor(verify_fail, Mode.Application).dryrun_one(
8585
args, compiler_version=5
8686
)
8787

@@ -117,7 +117,7 @@ def decompress():
117117
)
118118

119119
args = []
120-
app_result = PyTealDryRunExecutor(decompress, Mode.Application).dryrun(
120+
app_result = PyTealDryRunExecutor(decompress, Mode.Application).dryrun_one(
121121
args, compiler_version=5
122122
)
123123

@@ -158,7 +158,7 @@ def recover():
158158
)
159159

160160
args = []
161-
app_result = PyTealDryRunExecutor(recover, Mode.Application).dryrun(
161+
app_result = PyTealDryRunExecutor(recover, Mode.Application).dryrun_one(
162162
args, compiler_version=5
163163
)
164164

‎tests/integration/graviton_abi_test.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import random
22
import pytest
33

4-
from graviton.blackbox import DryRunInspector
4+
from graviton.inspector import DryRunInspector
55

66
import pyteal as pt
77
from pyteal.ast.subroutine import ABIReturnSubroutine
@@ -331,7 +331,7 @@ def pytuple_to_num(t):
331331
]
332332

333333
def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
334-
return p.dryrun_on_sequence(binary_inputs, compiler_version=version) # type: ignore
334+
return p.dryrun_sequence(binary_inputs, compiler_version=version) # type: ignore
335335

336336
# Binary:
337337
inspectors_subtract_slick = binary_dryrun(bbpt_subtract_slick)
@@ -343,7 +343,7 @@ def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
343343
inspectors_add = binary_dryrun(bbpt_add)
344344

345345
# Unary:
346-
inspectors_negate = bbpt_negate.dryrun_on_sequence(
346+
inspectors_negate = bbpt_negate.dryrun_sequence(
347347
unary_inputs, compiler_version=version # type: ignore
348348
)
349349

@@ -444,11 +444,11 @@ def pytuple_to_complex(tt):
444444

445445
# Binary:
446446
def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
447-
return p.dryrun_on_sequence(binary_inputs, compiler_version=version) # type: ignore
447+
return p.dryrun_sequence(binary_inputs, compiler_version=version) # type: ignore
448448

449449
# Unary:
450450
def unary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
451-
return p.dryrun_on_sequence(unary_inputs, compiler_version=version) # type: ignore
451+
return p.dryrun_sequence(unary_inputs, compiler_version=version) # type: ignore
452452

453453
inspectors_cplx_add = binary_dryrun(bbpt_cplx_add)
454454

@@ -523,7 +523,7 @@ def test_conditional_factorial(version: int):
523523
ptdre = PyTealDryRunExecutor(conditional_factorial, pt.Mode.Application)
524524
inputs = [(n,) for n in range(20)]
525525

526-
inspectors = ptdre.dryrun_on_sequence(inputs, compiler_version=version) # type: ignore
526+
inspectors = ptdre.dryrun_sequence(inputs, compiler_version=version) # type: ignore
527527
for i, args in enumerate(inputs):
528528
inspector = inspectors[i]
529529
n = args[0]
@@ -534,7 +534,7 @@ def test_conditional_factorial(version: int):
534534

535535
n = 21
536536
args = (n,)
537-
inspector = ptdre.dryrun(args)
537+
inspector = ptdre.dryrun_one(args)
538538
assert inspector.rejected(), inspector.report(
539539
args, f"FAILED: should have rejected for {n=}", row=n + 1
540540
)

0 commit comments

Comments
 (0)