diff --git a/pyteal/compiler/compiler.py b/pyteal/compiler/compiler.py index c465f00c4..ceec504d5 100644 --- a/pyteal/compiler/compiler.py +++ b/pyteal/compiler/compiler.py @@ -306,6 +306,7 @@ def __init__( *, version: int = DEFAULT_PROGRAM_VERSION, assemble_constants: bool = False, + assembly_type_track: bool = True, optimize: OptimizeOptions | None = None, ): """ @@ -322,12 +323,17 @@ def __init__( constants will be assembled in the most space-efficient way, so enabling this may reduce the compiled program's size. Enabling this option requires a minimum program version of 3. Defaults to `False`. + assembly_type_track (optional): When `True`, the compiler will produce a program with type + checking at assembly time (default behavior). When `False`, the compiler will turn off + type checking at assembly time. This is only useful if PyTeal is producing incorrect + TEAL code, or the assembler is producing incorrect type errors. Defaults to `True`. optimize (optional): `OptimizeOptions` that determine which optimizations will be applied. """ self.ast = ast self.mode = mode self.version = version self.assemble_constants = assemble_constants + self.assembly_type_track = assembly_type_track self.optimize: OptimizeOptions = optimize or OptimizeOptions() def compile( @@ -473,7 +479,11 @@ def _compile_impl( ) components = createConstantBlocks(components) - components = [TealPragma(self.version)] + components # T2PT0 + componentsPrefix: list[TealComponent] = [TealPragma(version=self.version)] + if not self.assembly_type_track: + componentsPrefix.append(TealPragma(type_track=False)) + + components = componentsPrefix + components # T2PT0 teal_chunks = [tl.assemble() for tl in components] teal_code = "\n".join(teal_chunks) @@ -536,6 +546,7 @@ def compileTeal( *, version: int = DEFAULT_PROGRAM_VERSION, assembleConstants: bool = False, + assembly_type_track: bool = True, optimize: OptimizeOptions | None = None, ) -> str: """Compile a PyTeal expression into TEAL assembly. @@ -551,6 +562,10 @@ def compileTeal( constants will be assembled in the most space-efficient way, so enabling this may reduce the compiled program's size. Enabling this option requires a minimum program version of 3. Defaults to false. + assembly_type_track (optional): When `True`, the compiler will produce a program with type + checking at assembly time (default behavior). When `False`, the compiler will turn off + type checking at assembly time. This is only useful if PyTeal is producing incorrect + TEAL code, or the assembler is producing incorrect type errors. Defaults to `True`. optimize (optional): OptimizeOptions that determine which optimizations will be applied. Returns: @@ -565,6 +580,7 @@ def compileTeal( mode, version=version, assemble_constants=assembleConstants, + assembly_type_track=assembly_type_track, optimize=optimize, )._compile_impl(with_sourcemap=False) return bundle.teal diff --git a/pyteal/compiler/compiler_test.py b/pyteal/compiler/compiler_test.py index 2f1f41881..c69f10394 100644 --- a/pyteal/compiler/compiler_test.py +++ b/pyteal/compiler/compiler_test.py @@ -217,6 +217,20 @@ def test_compile_version_6(): assert actual == expected +def test_compile_no_assembly_type_tracking(): + expr = pt.Int(1) + expected = """ +#pragma version 6 +#pragma typetrack false +int 1 +return +""".strip() + actual = pt.compileTeal( + expr, pt.Mode.Signature, version=6, assembly_type_track=False + ) + assert actual == expected + + def test_slot_load_before_store(): program = pt.AssetHolding.balance(pt.Int(0), pt.Int(0)).value() with pytest.raises(pt.TealInternalError): diff --git a/pyteal/ir/tealpragma.py b/pyteal/ir/tealpragma.py index 4a0ff5d46..b8348c792 100644 --- a/pyteal/ir/tealpragma.py +++ b/pyteal/ir/tealpragma.py @@ -2,15 +2,47 @@ class TealPragma(TealComponent): - def __init__(self, version: int): + _name: str + _value: str | int + + def __init__(self, version: int | None = None, *, type_track: bool | None = None): + """Creates an assembler pragma statement. + + Only one of the arguments should be set. + + Args: + version (optional): Sets the program version. + type_track (optional): Configures assembler type tracking. + """ super().__init__(None) - self.version = version + + if len([x for x in [version, type_track] if x is not None]) != 1: + raise ValueError("Exactly one of version or type_track must be set") + + if version is not None: + self._name = "version" + self._value = version + elif type_track is not None: + self._name = "typetrack" + self._value = "true" if type_track else "false" + else: + # Shouldn't happen, just to satisfy type checker + raise ValueError("Empty pragma statement") def assemble(self) -> str: - return f"#pragma version {self.version}" + return f"#pragma {self._name} {self._value}" def __repr__(self) -> str: - return f"TealPragma({self.version})" + match self._name: + case "version": + name = "version" + value = self._value + case "typetrack": + name = "type_track" + value = self._value == "true" + case _: + raise ValueError(f"Unknown pragma name: {self._name}") + return f"TealPragma({name}={value})" def __hash__(self) -> int: return hash(repr(self)) @@ -18,7 +50,7 @@ def __hash__(self) -> int: def __eq__(self, other: object) -> bool: if not isinstance(other, TealPragma): return False - return self.version == other.version + return self._name == other._name and self._value == other._value TealPragma.__module__ = "pyteal" diff --git a/pyteal/ir/tealpragma_test.py b/pyteal/ir/tealpragma_test.py new file mode 100644 index 000000000..251b7cb76 --- /dev/null +++ b/pyteal/ir/tealpragma_test.py @@ -0,0 +1,30 @@ +import pytest +import pyteal as pt + + +def test_version(): + for i in range(10): + version_pragma = pt.TealPragma(version=i) + assert version_pragma._name == "version" + assert version_pragma._value == i + assert version_pragma.assemble() == f"#pragma version {i}" + assert repr(version_pragma) == f"TealPragma(version={i})" + + +def test_type_track(): + for value in (True, False): + type_track_pragma = pt.TealPragma(type_track=value) + assert type_track_pragma._name == "typetrack" + assert type_track_pragma._value == str(value).lower() + assert type_track_pragma.assemble() == f"#pragma typetrack {str(value).lower()}" + assert repr(type_track_pragma) == f"TealPragma(type_track={value})" + + +def test_empty(): + with pytest.raises(ValueError): + pt.TealPragma() + + +def test_both(): + with pytest.raises(ValueError): + pt.TealPragma(version=1, type_track=True) diff --git a/tests/blackbox.py b/tests/blackbox.py index 0b0aafc9d..65c6a6074 100644 --- a/tests/blackbox.py +++ b/tests/blackbox.py @@ -434,12 +434,13 @@ def approval(): return approval - def compile(self, version: int, assemble_constants: bool = False) -> str: + def compile(self, version: int, *, assemble_constants: bool = False, assembly_type_track: bool = True) -> str: return compileTeal( self.program(), self.mode, version=version, assembleConstants=assemble_constants, + assembly_type_track=assembly_type_track, ) def executor(self, compiler_version: int = 6) -> DryRunExecutor: diff --git a/tests/integration/graviton_test.py b/tests/integration/graviton_test.py index 988db713a..3340fa0bd 100644 --- a/tests/integration/graviton_test.py +++ b/tests/integration/graviton_test.py @@ -31,11 +31,11 @@ def wrap_compile_and_save( - subr, mode, version, assemble_constants, test_name, case_name + subr, mode, version, assemble_constants, test_name, case_name, *, assembly_type_track: bool = True, ): is_app = mode == pt.Mode.Application - teal = PyTealDryRunExecutor(subr, mode).compile(version, assemble_constants) + teal = PyTealDryRunExecutor(subr, mode).compile(version, assemble_constants=assemble_constants, assembly_type_track=assembly_type_track) tealfile = f'{"app" if is_app else "lsig"}_{case_name}_v{version}.teal' tealdir = GENERATED / test_name @@ -481,7 +481,9 @@ def blackbox_test_runner( # 1. Compile to TEAL teal, _, tealfile = wrap_compile_and_save( - subr, mode, version, assemble_constants, "blackbox", case_name + subr, mode, version, assemble_constants, "blackbox", case_name, + # Temporarily disabling until https://github.com/algorand/go-algorand/pull/5884 is released + assembly_type_track=False, ) # Fail fast in case algod is not configured: