Refactor Uint8.set to make data size a first class concept

pyteal/ast/abi/sized.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from typing import Union
+from .. import Expr
+from ..assert_ import Assert
+from pyteal.ast import Int, Seq, ScratchVar
+
+
+@dataclass(frozen=True)
+class SizedExpr:
+    """Provides a wrapper class to tag `Expr`s with known ABI data sizes.
+
+    Knowing the data size removes the need for a runtime `Assert` confirming it fits into the ScratchVar."""
+
+    underlying: Expr
+
+    @staticmethod
+    def store_into(
+        u: Union[Expr, "SizedExpr"], s: ScratchVar, supported_bit_size: int
+    ) -> Expr:
+        return (
+            s.store(u.underlying)
+            if isinstance(u, SizedExpr)
+            else Seq(s.store(u), Assert(s.load() < Int(2 ** supported_bit_size)))
+        )

pyteal/ast/abi/uint.py

@@ -1,6 +1,7 @@
 from typing import Union, cast
 from abc import abstractmethod
 
+from .sized import SizedExpr
 from ...types import TealType
 from ...errors import TealInputError
 from ..expr import Expr
@@ -56,36 +57,34 @@ def new_instance(self) -> "Uint8":
         return Uint8()
 
     def set(self, value: Union[int, Expr, "Uint8", "Byte"]) -> Expr:
-        checked = False
-        if type(value) is int:
-            if value >= 2 ** self.bit_size:
-                raise TealInputError(
-                    "Value exceeds {} maximum: {}".format(
-                        self.__class__.__name__, value
+        def resolve_storage_value() -> Union[Expr, SizedExpr]:
+            if type(value) is int:
+                if value >= 2 ** self.bit_size:
+                    raise TealInputError(
+                        "Value exceeds {} maximum: {}".format(
+                            self.__class__.__name__, value
+                        )
                     )
-                )
-            value = Int(value)
-            checked = True
-
-        if type(value) is Uint8 or type(value) is Byte:
-            value = value.get()
-            checked = True
+                return SizedExpr(Int(value))
 
-        if checked:
-            return self.stored_value.store(cast(Expr, value))
+            elif isinstance(value, Uint8):
+                return SizedExpr(value.get())
+            elif isinstance(value, Expr):
+                # There's insufficient evidence to attempt resolution.
+                return cast(Expr, value)
+            else:
+                raise TypeError(f"Unsupported type: {type(value)}")
 
-        return Seq(
-            self.stored_value.store(cast(Expr, value)),
-            Assert(self.stored_value.load() < Int(2 ** self.bit_size)),
-        )
+        r = resolve_storage_value()
+        return SizedExpr.store_into(r, self.stored_value, self.bit_size)
 
     def decode(
         self,
         encoded: Expr,
         *,
         startIndex: Expr = None,
         endIndex: Expr = None,
-        length: Expr = None
+        length: Expr = None,
     ) -> Expr:
         if startIndex is None:
             startIndex = Int(0)
@@ -147,7 +146,7 @@ def decode(
         *,
         startIndex: Expr = None,
         endIndex: Expr = None,
-        length: Expr = None
+        length: Expr = None,
     ) -> Expr:
         if startIndex is None:
             startIndex = Int(0)
@@ -195,7 +194,7 @@ def decode(
         *,
         startIndex: Expr = None,
         endIndex: Expr = None,
-        length: Expr = None
+        length: Expr = None,
     ) -> Expr:
         if startIndex is None:
             startIndex = Int(0)
@@ -228,7 +227,7 @@ def decode(
         *,
         startIndex: Expr = None,
         endIndex: Expr = None,
-        length: Expr = None
+        length: Expr = None,
     ) -> Expr:
         if startIndex is None:
             if endIndex is None and length is None:

requirements.txt

@@ -3,3 +3,4 @@ mypy==0.931
 pytest
 pytest-timeout
 py-algorand-sdk
+types-dataclasses