Skip to content

Commit

Permalink
make UInt64 backed enums subtypes of UInt64
Browse files Browse the repository at this point in the history
  • Loading branch information
daniel-makerx committed Jul 24, 2024
1 parent 75cc43d commit 3b95140
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 16 deletions.
Empty file added examples/box/__init__.py
Empty file.
20 changes: 20 additions & 0 deletions examples/box/contract.py
Original file line number Diff line number Diff line change
@@ -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
26 changes: 26 additions & 0 deletions examples/box/test_contract.py
Original file line number Diff line number Diff line change
@@ -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)
50 changes: 35 additions & 15 deletions src/algopy_testing/enums.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down
6 changes: 5 additions & 1 deletion src/algopy_testing/models/box.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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]
Expand Down

0 comments on commit 3b95140

Please sign in to comment.