Use named class methods to specify value type

pyteal/ast/jsonref.py

@@ -2,7 +2,7 @@
 from enum import Enum
 
 from pyteal.types import TealType, require_type
-from pyteal.errors import TealTypeError, verifyFieldVersion, verifyTealVersion
+from pyteal.errors import verifyFieldVersion, verifyTealVersion
 from pyteal.ir import TealOp, Op, TealBlock
 from pyteal.ast.expr import Expr
 from pyteal.ast.leafexpr import LeafExpr
@@ -14,9 +14,9 @@
 class JsonRefType(Enum):
     # fmt: off
     #              id |    name   |      type      | min version
-    json_string = (0, "JSONString", TealType.bytes,  7)
-    json_uint64 = (1, "JSONUint64", TealType.uint64, 7)
-    json_object = (2, "JSONObject", TealType.bytes,  7)
+    string = (0, "JSONString", TealType.bytes,  7)
+    uint64 = (1, "JSONUint64", TealType.uint64, 7)
+    object = (2, "JSONObject", TealType.bytes,  7)
     # fmt: on
 
     def __init__(self, id: int, name: str, type: TealType, min_version: int) -> None:
@@ -33,20 +33,14 @@ def type_of(self) -> TealType:
 
 
 class JsonRef(LeafExpr):
-    """An expression that accesses a key from a supported utf-8 encoded json object.
+    """An expression that accesses the value associated with a given key from a supported utf-8 encoded json object.
 
     The json object must satisfy a [particular specification](https://github.com/algorand/go-algorand/blob/master/data/transactions/logic/jsonspec.md).
     """
 
-    json_string = JsonRefType.json_string
-    json_uint64 = JsonRefType.json_uint64
-    json_object = JsonRefType.json_object
-
     def __init__(self, type: JsonRefType, json_obj: Expr, key: Expr) -> None:
         super().__init__()
 
-        if not isinstance(type, JsonRefType):
-            raise TealTypeError(type, JsonRefType)
         self.type = type
 
         require_type(json_obj, TealType.bytes)
@@ -73,5 +67,41 @@ def __str__(self):
     def type_of(self):
         return self.type.type_of()
 
+    @classmethod
+    def as_string(cls, json_obj: Expr, key: Expr) -> "JsonRef":
+        """Access the value of a given key as a string.
+
+        Refer to the `JsonRef` class documentation for valid json specification.
+
+        Args:
+            json_obj: The utf-8 encoded json object.
+            key: The key to access in the json object.
+        """
+        return cls(JsonRefType.string, json_obj, key)
+
+    @classmethod
+    def as_uint64(cls, json_obj: Expr, key: Expr) -> "JsonRef":
+        """Access the value of a given key as a uint64.
+
+        Refer to the `JsonRef` class documentation for valid json specification.
+
+        Args:
+            json_obj: The utf-8 encoded json object.
+            key: The key to access in the json object.
+        """
+        return cls(JsonRefType.uint64, json_obj, key)
+
+    @classmethod
+    def as_object(cls, json_obj: Expr, key: Expr) -> "JsonRef":
+        """Access the value of a given key as a json object.
+
+        Refer to the `JsonRef` class documentation for valid json specification.
+
+        Args:
+            json_obj: The utf-8 encoded json object.
+            key: The key to access in the json object.
+        """
+        return cls(JsonRefType.object, json_obj, key)
+
 
 JsonRef.__module__ = "pyteal"

pyteal/ast/jsonref_test.py

@@ -8,7 +8,7 @@
 
 def test_json_string():
     args = [pt.Bytes('{"foo":"bar"}'), pt.Bytes("foo")]
-    expr = pt.JsonRef(pt.JsonRef.json_string, *args)
+    expr = pt.JsonRef.as_string(*args)
     assert expr.type_of() == pt.TealType.bytes
 
     expected = pt.TealSimpleBlock(
@@ -32,7 +32,7 @@ def test_json_string():
 
 def test_json_uint64():
     args = [pt.Bytes('{"foo":123456789}'), pt.Bytes("foo")]
-    expr = pt.JsonRef(pt.JsonRef.json_uint64, *args)
+    expr = pt.JsonRef.as_uint64(*args)
     assert expr.type_of() == pt.TealType.uint64
 
     expected = pt.TealSimpleBlock(
@@ -56,7 +56,7 @@ def test_json_uint64():
 
 def test_json_object():
     args = [pt.Bytes('{"foo":{"key": "value"}}'), pt.Bytes("foo")]
-    expr = pt.JsonRef(pt.JsonRef.json_object, *args)
+    expr = pt.JsonRef.as_object(*args)
     assert expr.type_of() == pt.TealType.bytes
 
     expected = pt.TealSimpleBlock(
@@ -80,10 +80,7 @@ def test_json_object():
 
 def test_json_ref_invalid():
     with pytest.raises(pt.TealTypeError):
-        pt.JsonRef(pt.Bytes("my string"), pt.Bytes("a"), pt.Bytes("a"))
+        pt.JsonRef.as_object(pt.Int(0), pt.Bytes("a"))
 
     with pytest.raises(pt.TealTypeError):
-        pt.JsonRef(pt.JsonRef.json_object, pt.Int(0), pt.Bytes("a"))
-
-    with pytest.raises(pt.TealTypeError):
-        pt.JsonRef(pt.JsonRef.json_string, pt.Bytes("a"), pt.Int(0))
+        pt.JsonRef.as_string(pt.Bytes("a"), pt.Int(0))