Optionally refactor handling of array_length int subclasses in Static…

…ArrayTypeSpec (#347)
algorand · May 17, 2022 · 00971c4 · 00971c4
1 parent 1ee6151
commit 00971c4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py
@@ -22,7 +22,7 @@ class AddressLength(IntEnum):
 
 class AddressTypeSpec(StaticArrayTypeSpec):
     def __init__(self) -> None:
-        super().__init__(ByteTypeSpec(), int(AddressLength.Bytes))
+        super().__init__(ByteTypeSpec(), AddressLength.Bytes)
 
     def new_instance(self) -> "Address":
         return Address()
@@ -67,7 +67,7 @@ def set(
             case ComputedValue():
                 pts = value.produced_type_spec()
                 if pts == AddressTypeSpec() or pts == StaticArrayTypeSpec(
-                    ByteTypeSpec(), int(AddressLength.Bytes)
+                    ByteTypeSpec(), AddressLength.Bytes
                 ):
                     return value.store_into(self)
 
@@ -78,7 +78,7 @@ def set(
                 if (
                     value.type_spec() == AddressTypeSpec()
                     or value.type_spec()
-                    == StaticArrayTypeSpec(ByteTypeSpec(), int(AddressLength.Bytes))
+                    == StaticArrayTypeSpec(ByteTypeSpec(), AddressLength.Bytes)
                 ):
                     return self.stored_value.store(value.stored_value.load())
 

diff --git a/pyteal/ast/abi/array_static.py b/pyteal/ast/abi/array_static.py
@@ -18,7 +18,9 @@ def __init__(self, value_type_spec: TypeSpec, array_length: int) -> None:
         super().__init__(value_type_spec)
         if not isinstance(array_length, int) or array_length < 0:
             raise TypeError(f"Unsupported StaticArray length: {array_length}")
-        self.array_length: Final = array_length
+
+        # Casts to `int` to handle downstream usage where value is a subclass of int like `IntEnum`.
+        self.array_length: Final = int(array_length)
 
     def new_instance(self) -> "StaticArray[T, N]":
         return StaticArray(self)