make UInt64 backed enums subtypes of UInt64

Author: daniel-makerx

examples/box/__init__.py

@@ -0,0 +1,20 @@
+from algopy import ARC4Contract, Box, OnCompleteAction, TransactionType, arc4, op
+
+
+class BoxContract(ARC4Contract):
+
+    def __init__(self) -> None:
+        self.oca = Box(OnCompleteAction)
+        self.txn = Box(TransactionType)
+
+    @arc4.abimethod()
+    def store_enums(self) -> None:
+        self.oca.value = OnCompleteAction.OptIn
+        self.txn.value = TransactionType.ApplicationCall
+
+    @arc4.abimethod()
+    def read_enums(self) -> tuple[OnCompleteAction, TransactionType]:
+        assert op.Box.get(b"oca")[0] == op.itob(self.oca.value)
+        assert op.Box.get(b"txn")[0] == op.itob(self.txn.value)
+
+        return self.oca.value, self.txn.value

examples/box/contract.py

@@ -0,0 +1,26 @@
+from collections.abc import Generator
+
+import pytest
+from algopy import op
+from algopy_testing import AlgopyTestContext, algopy_testing_context
+
+from .contract import BoxContract
+
+
+@pytest.fixture()
+def context() -> Generator[AlgopyTestContext, None, None]:
+    with algopy_testing_context() as ctx:
+        yield ctx
+
+
+def test_enums(context: AlgopyTestContext) -> None:
+    # Arrange
+    contract = BoxContract()
+
+    # Act
+    contract.store_enums()
+    oca, txn = contract.read_enums()
+
+    # Assert
+    assert context.get_box(b"oca") == op.itob(oca)
+    assert context.get_box(b"txn") == op.itob(txn)

examples/box/test_contract.py

@@ -1,22 +1,42 @@
-from enum import Enum, IntEnum, StrEnum
+from __future__ import annotations
 
+from enum import Enum, StrEnum
 
-class OnCompleteAction(Enum):
-    NoOp = 0
-    OptIn = 1
-    CloseOut = 2
-    ClearState = 3
-    UpdateApplication = 4
-    DeleteApplication = 5
+from algopy_testing.primitives import UInt64
 
 
-class TransactionType(IntEnum):
-    Payment = 0
-    KeyRegistration = 1
-    AssetConfig = 2
-    AssetTransfer = 3
-    AssetFreeze = 4
-    ApplicationCall = 5
+class OnCompleteAction(UInt64):
+    NoOp: OnCompleteAction
+    OptIn: OnCompleteAction
+    CloseOut: OnCompleteAction
+    ClearState: OnCompleteAction
+    UpdateApplication: OnCompleteAction
+    DeleteApplication: OnCompleteAction
+
+
+OnCompleteAction.NoOp = OnCompleteAction(0)
+OnCompleteAction.OptIn = OnCompleteAction(1)
+OnCompleteAction.CloseOut = OnCompleteAction(2)
+OnCompleteAction.ClearState = OnCompleteAction(3)
+OnCompleteAction.UpdateApplication = OnCompleteAction(4)
+OnCompleteAction.DeleteApplication = OnCompleteAction(5)
+
+
+class TransactionType(UInt64):
+    Payment: TransactionType
+    KeyRegistration: TransactionType
+    AssetConfig: TransactionType
+    AssetTransfer: TransactionType
+    AssetFreeze: TransactionType
+    ApplicationCall: TransactionType
+
+
+TransactionType.Payment = TransactionType(0)
+TransactionType.KeyRegistration = TransactionType(1)
+TransactionType.AssetConfig = TransactionType(2)
+TransactionType.AssetTransfer = TransactionType(3)
+TransactionType.AssetFreeze = TransactionType(4)
+TransactionType.ApplicationCall = TransactionType(5)
 
 
 class ECDSA(Enum):

src/algopy_testing/enums.py

@@ -450,7 +450,7 @@ def _full_key(self, key: _TKey) -> algopy.Bytes:
         return self.key_prefix + _cast_to_bytes(key)
 
 
-def _cast_to_value_type(t: type[_TValue], value: bytes) -> _TValue:
+def _cast_to_value_type(t: type[_TValue], value: bytes) -> _TValue:  # noqa: PLR0911
     """
     assuming _TValue to be one of the followings:
         - bool,
@@ -474,6 +474,10 @@ def _cast_to_value_type(t: type[_TValue], value: bytes) -> _TValue:
         return algopy.Bytes(value)  # type: ignore[return-value]
     elif t is algopy.UInt64:
         return algopy.op.btoi(value)  # type: ignore[return-value]
+    elif t is algopy.OnCompleteAction:
+        return algopy.OnCompleteAction(algopy.op.btoi(value).value)  # type: ignore[return-value]
+    elif t is algopy.TransactionType:
+        return algopy.TransactionType(algopy.op.btoi(value).value)  # type: ignore[return-value]
     elif t is algopy.Asset:
         asset_id = algopy.op.btoi(value)
         return context.get_asset(asset_id)  # type: ignore[return-value]