Refactor graviton usage (#628)

requirements.txt

@@ -1,7 +1,7 @@
 black==22.3.0
 flake8==5.0.4
 flake8-tidy-imports==4.6.0
-graviton@git+https://github.com/algorand/graviton@v0.7.1
+graviton@git+https://github.com/algorand/graviton@v0.8.0
 mypy==0.991
 pytest==7.2.0
 pytest-cov==3.0.0

tests/blackbox.py

@@ -1,11 +1,14 @@
-from typing import Any, Callable, Generic, Optional, Sequence, TypeVar, cast
-from dataclasses import dataclass
+from typing import Any, Callable, Sequence, cast
 
 import algosdk.abi
 from algosdk.v2client import algod
 
 from graviton import blackbox
-from graviton.blackbox import DryRunInspector, DryRunExecutor
+from graviton.blackbox import (
+    DryRunInspector,
+    DryRunExecutor,
+    DryRunTransactionParams as TxParams,
+)
 from graviton.models import PyTypes
 
 from pyteal.ast.subroutine import OutputKwArgInfo
@@ -138,37 +141,13 @@ def decorator_blackbox(func: SubroutineFnWrapper | ABIReturnSubroutine):
 # ---- API ---- #
 
 
-Output = TypeVar("Output")
-Lazy = Callable[[], Output]
-
-
-@dataclass(frozen=True)
-class _MatchMode(Generic[Output]):
-    runner: Optional["PyTealDryRunExecutor"]
-    app_case: Lazy
-    signature_case: Lazy
-    trace: Any = None
-
-    def __post_init__(self):
-        if self.runner and self.trace:
-            self.runner.add_trace(self.trace)
-
-    def __call__(self, mode: Mode, *args, **kwargs) -> Output:
-        match mode:
-            case Mode.Application:
-                return self.app_case()
-            case Mode.Signature:
-                return self.signature_case()
-            case _:
-                raise Exception(f"Unknown mode {mode} of type {type(mode)}")
-
-
 def mode_to_execution_mode(mode: Mode) -> blackbox.ExecutionMode:
-    return _MatchMode(
-        runner=None,
-        app_case=lambda: blackbox.ExecutionMode.Application,
-        signature_case=lambda: blackbox.ExecutionMode.Signature,
-    )(mode)
+    if mode == Mode.Application:
+        return blackbox.ExecutionMode.Application
+    if mode == Mode.Signature:
+        return blackbox.ExecutionMode.Signature
+
+    raise ValueError(f"Can't handle {mode=}")
 
 
 class PyTealDryRunExecutor:
@@ -182,9 +161,11 @@ def __init__(self, subr: BlackboxWrapper, mode: Mode):
             mode: type of program to produce: logic sig (Mode.Signature) or app (Mode.Application)
         """
         input_types = subr.input_types
-        assert (
-            input_types is not None
-        ), "please provide input_types in your @Subroutine or @ABIReturnSubroutine annotation (this is crucial for generating proper end-to-end testable PyTeal)"
+        assert input_types is not None, (
+            "please provide input_types in your @Subroutine or @ABIReturnSubroutine "
+            "annotation. "
+            "(this is crucial for generating proper end-to-end testable PyTeal)"
+        )
 
         self.subr, self.mode, self.input_types = subr, mode, input_types
         match subr.subroutine:
@@ -417,68 +398,45 @@ def approval():
         return approval
 
     def compile(self, version: int, assemble_constants: bool = False) -> str:
-        return _MatchMode(
-            runner=self,
-            app_case=lambda: compileTeal(
-                self.program(),
-                self.mode,
-                version=version,
-                assembleConstants=assemble_constants,
-            ),
-            signature_case=lambda: compileTeal(
-                self.program(),
-                self.mode,
-                version=version,
-                assembleConstants=assemble_constants,
-            ),
-        )(self.mode)
-
-    def dryrun_on_sequence(
+        return compileTeal(
+            self.program(),
+            self.mode,
+            version=version,
+            assembleConstants=assemble_constants,
+        )
+
+    def executor(self, compiler_version: int = 6) -> DryRunExecutor:
+        return DryRunExecutor(
+            algod=algod_with_assertion(),
+            mode=mode_to_execution_mode(self.mode),
+            teal=self.compile(compiler_version),
+            abi_method_signature=self.abi_method_signature(),
+            omit_method_selector=True,
+        )
+
+    def dryrun_sequence(
         self,
         inputs: list[Sequence[PyTypes]],
+        *,
         compiler_version=6,
+        txn_params: TxParams | None = None,
+        verbose: bool = False,
     ) -> list[DryRunInspector]:
-        teal = self.compile(compiler_version)
-        return _MatchMode(
-            self,
-            app_case=lambda: DryRunExecutor.dryrun_app_on_sequence(
-                algod=algod_with_assertion(),
-                teal=teal,
-                inputs=inputs,
-                abi_method_signature=self.abi_method_signature(),
-                omit_method_selector=True,
+        return cast(
+            list,
+            self.executor(compiler_version).run_sequence(
+                inputs, txn_params=txn_params, verbose=verbose
             ),
-            signature_case=lambda: DryRunExecutor.dryrun_logicsig_on_sequence(
-                algod=algod_with_assertion(),
-                teal=teal,
-                inputs=inputs,
-                abi_method_signature=self.abi_method_signature(),
-                omit_method_selector=True,
-            ),
-            trace=teal,
-        )(self.mode)
+        )
 
-    def dryrun(
+    def dryrun_one(
         self,
         args: Sequence[bytes | str | int],
+        *,
         compiler_version=6,
+        txn_params: TxParams | None = None,
+        verbose: bool = False,
     ) -> DryRunInspector:
-        teal = self.compile(compiler_version)
-        return _MatchMode(
-            self,
-            app_case=lambda: DryRunExecutor.dryrun_app(
-                algod_with_assertion(),
-                teal,
-                args,
-                abi_method_signature=self.abi_method_signature(),
-                omit_method_selector=True,
-            ),
-            signature_case=lambda: DryRunExecutor.dryrun_logicsig(
-                algod_with_assertion(),
-                teal,
-                args,
-                abi_method_signature=self.abi_method_signature(),
-                omit_method_selector=True,
-            ),
-            trace=teal,
-        )(self.mode)
+        return self.executor(compiler_version).run_one(
+            args, txn_params=txn_params, verbose=verbose
+        )

tests/integration/abi_roundtrip_test.py

@@ -257,7 +257,7 @@ def test_roundtrip(abi_type):
     args = (rand_abi_instance,)
 
     def dryrun_roundtrip(version: int):
-        inspector = roundtripper.dryrun(args, compiler_version=version)
+        inspector = roundtripper.dryrun_one(args, compiler_version=version)
 
         cost = inspector.cost()
         passed = inspector.passed()

tests/integration/ecdsa_test.py

@@ -46,7 +46,7 @@ def verify():
         )
 
     args = []
-    app_result = PyTealDryRunExecutor(verify, Mode.Application).dryrun(
+    app_result = PyTealDryRunExecutor(verify, Mode.Application).dryrun_one(
         args, compiler_version=5
     )
 
@@ -81,7 +81,7 @@ def verify_fail():
         )
 
     args = []
-    app_result = PyTealDryRunExecutor(verify_fail, Mode.Application).dryrun(
+    app_result = PyTealDryRunExecutor(verify_fail, Mode.Application).dryrun_one(
         args, compiler_version=5
     )
 
@@ -117,7 +117,7 @@ def decompress():
         )
 
     args = []
-    app_result = PyTealDryRunExecutor(decompress, Mode.Application).dryrun(
+    app_result = PyTealDryRunExecutor(decompress, Mode.Application).dryrun_one(
         args, compiler_version=5
     )
 
@@ -158,7 +158,7 @@ def recover():
         )
 
     args = []
-    app_result = PyTealDryRunExecutor(recover, Mode.Application).dryrun(
+    app_result = PyTealDryRunExecutor(recover, Mode.Application).dryrun_one(
         args, compiler_version=5
     )
 

tests/integration/graviton_abi_test.py

@@ -1,7 +1,7 @@
 import random
 import pytest
 
-from graviton.blackbox import DryRunInspector
+from graviton.inspector import DryRunInspector
 
 import pyteal as pt
 from pyteal.ast.subroutine import ABIReturnSubroutine
@@ -331,7 +331,7 @@ def pytuple_to_num(t):
     ]
 
     def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
-        return p.dryrun_on_sequence(binary_inputs, compiler_version=version)  # type: ignore
+        return p.dryrun_sequence(binary_inputs, compiler_version=version)  # type: ignore
 
     # Binary:
     inspectors_subtract_slick = binary_dryrun(bbpt_subtract_slick)
@@ -343,7 +343,7 @@ def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
     inspectors_add = binary_dryrun(bbpt_add)
 
     # Unary:
-    inspectors_negate = bbpt_negate.dryrun_on_sequence(
+    inspectors_negate = bbpt_negate.dryrun_sequence(
         unary_inputs, compiler_version=version  # type: ignore
     )
 
@@ -444,11 +444,11 @@ def pytuple_to_complex(tt):
 
     # Binary:
     def binary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
-        return p.dryrun_on_sequence(binary_inputs, compiler_version=version)  # type: ignore
+        return p.dryrun_sequence(binary_inputs, compiler_version=version)  # type: ignore
 
     # Unary:
     def unary_dryrun(p: PyTealDryRunExecutor) -> list[DryRunInspector]:
-        return p.dryrun_on_sequence(unary_inputs, compiler_version=version)  # type: ignore
+        return p.dryrun_sequence(unary_inputs, compiler_version=version)  # type: ignore
 
     inspectors_cplx_add = binary_dryrun(bbpt_cplx_add)
 
@@ -523,7 +523,7 @@ def test_conditional_factorial(version: int):
     ptdre = PyTealDryRunExecutor(conditional_factorial, pt.Mode.Application)
     inputs = [(n,) for n in range(20)]
 
-    inspectors = ptdre.dryrun_on_sequence(inputs, compiler_version=version)  # type: ignore
+    inspectors = ptdre.dryrun_sequence(inputs, compiler_version=version)  # type: ignore
     for i, args in enumerate(inputs):
         inspector = inspectors[i]
         n = args[0]
@@ -534,7 +534,7 @@ def test_conditional_factorial(version: int):
 
     n = 21
     args = (n,)
-    inspector = ptdre.dryrun(args)
+    inspector = ptdre.dryrun_one(args)
     assert inspector.rejected(), inspector.report(
         args, f"FAILED: should have rejected for {n=}", row=n + 1
     )