Add JsonRef (#417)

* Add JsonRef

* Use named class methods to specify value type

* Remove unnecessary ignore Expr equality context

* Fix docstring link

pyteal/__init__.pyi

@@ -108,6 +108,7 @@ __all__ = [
     "InnerTxnGroup",
     "Int",
     "Itob",
+    "JsonRef",
     "Keccak256",
     "LabelReference",
     "Le",

pyteal/ast/__init__.py

@@ -104,6 +104,7 @@
 from pyteal.ast.ternaryexpr import Divw, Ed25519Verify, SetBit, SetByte
 from pyteal.ast.substring import Substring, Extract, Suffix
 from pyteal.ast.replace import Replace
+from pyteal.ast.jsonref import JsonRef
 
 # more ops
 from pyteal.ast.naryexpr import NaryExpr, And, Or, Concat
@@ -285,4 +286,5 @@
     "EcdsaVerify",
     "EcdsaDecompress",
     "EcdsaRecover",
+    "JsonRef",
 ]

pyteal/ast/jsonref.py

@@ -0,0 +1,107 @@
+from typing import TYPE_CHECKING
+from enum import Enum
+
+from pyteal.types import TealType, require_type
+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
+
+if TYPE_CHECKING:
+    from pyteal.compiler import CompileOptions
+
+
+class JsonRefType(Enum):
+    # fmt: off
+    #              id |    name   |      type      | min version
+    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:
+        self.id = id
+        self.arg_name = name
+        self.ret_type = type
+        self.min_version = min_version
+
+    def type_of(self) -> TealType:
+        return self.ret_type
+
+
+JsonRefType.__module__ = "pyteal"
+
+
+class JsonRef(LeafExpr):
+    """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>`_.
+    """
+
+    def __init__(self, type: JsonRefType, json_obj: Expr, key: Expr) -> None:
+        super().__init__()
+
+        self.type = type
+
+        require_type(json_obj, TealType.bytes)
+        self.json_obj = json_obj
+
+        require_type(key, TealType.bytes)
+        self.key = key
+
+    def __teal__(self, options: "CompileOptions"):
+        verifyTealVersion(
+            Op.json_ref.min_version,
+            options.version,
+            "TEAL version too low to use op json_ref",
+        )
+
+        verifyFieldVersion(self.type.arg_name, self.type.min_version, options.version)
+
+        op = TealOp(self, Op.json_ref, self.type.arg_name)
+        return TealBlock.FromOp(options, op, self.json_obj, self.key)
+
+    def __str__(self):
+        return "(JsonRef {})".format(self.type.arg_name)
+
+    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

@@ -0,0 +1,83 @@
+import pytest
+
+import pyteal as pt
+
+teal6Options = pt.CompileOptions(version=6)
+teal7Options = pt.CompileOptions(version=7)
+
+
+def test_json_string():
+    args = [pt.Bytes('{"foo":"bar"}'), pt.Bytes("foo")]
+    expr = pt.JsonRef.as_string(*args)
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{\\"foo\\":\\"bar\\"}"'),
+            pt.TealOp(args[1], pt.Op.byte, '"foo"'),
+            pt.TealOp(expr, pt.Op.json_ref, "JSONString"),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+def test_json_uint64():
+    args = [pt.Bytes('{"foo":123456789}'), pt.Bytes("foo")]
+    expr = pt.JsonRef.as_uint64(*args)
+    assert expr.type_of() == pt.TealType.uint64
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{\\"foo\\":123456789}"'),
+            pt.TealOp(args[1], pt.Op.byte, '"foo"'),
+            pt.TealOp(expr, pt.Op.json_ref, "JSONUint64"),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+def test_json_object():
+    args = [pt.Bytes('{"foo":{"key": "value"}}'), pt.Bytes("foo")]
+    expr = pt.JsonRef.as_object(*args)
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{\\"foo\\":{\\"key\\": \\"value\\"}}"'),
+            pt.TealOp(args[1], pt.Op.byte, '"foo"'),
+            pt.TealOp(expr, pt.Op.json_ref, "JSONObject"),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+def test_json_ref_invalid():
+    with pytest.raises(pt.TealTypeError):
+        pt.JsonRef.as_object(pt.Int(0), pt.Bytes("a"))
+
+    with pytest.raises(pt.TealTypeError):
+        pt.JsonRef.as_string(pt.Bytes("a"), pt.Int(0))

pyteal/ir/ops.py

@@ -181,6 +181,7 @@ def min_version(self) -> int:
     acct_params_get     = OpType("acct_params_get",     Mode.Application,                  6)
     replace2            = OpType("replace2",            Mode.Signature | Mode.Application, 7)
     replace3            = OpType("replace3",            Mode.Signature | Mode.Application, 7)
+    json_ref            = OpType("json_ref",            Mode.Signature | Mode.Application, 7)
     block               = OpType("block",               Mode.Signature | Mode.Application, 7)
     # fmt: on