Skip to content

Commit 39a11c2

Browse files
jdtzmnmichaeldiamant
jdtzmn
and
michaeldiamant
authored
Add Replace (#413)
* Add Replace * Remove replace auto-import * Use scripts/generate_init.py * Add more tests to replace, substring, and extract (#1) Co-authored-by: Michael Diamant <michaeldiamant@users.noreply.github.com>
1 parent 749fb8c commit 39a11c2

File tree

7 files changed

+276
-6
lines changed

7 files changed

+276
-6
lines changed

‎pyteal/__init__.pyi‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ __all__ = [
138138
"Or",
139139
"Pop",
140140
"Reject",
141+
"Replace",
141142
"Return",
142143
"ScratchIndex",
143144
"ScratchLoad",

‎pyteal/ast/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@
102102
# ternary ops
103103
from pyteal.ast.ternaryexpr import Divw, Ed25519Verify, SetBit, SetByte
104104
from pyteal.ast.substring import Substring, Extract, Suffix
105+
from pyteal.ast.replace import Replace
105106

106107
# more ops
107108
from pyteal.ast.naryexpr import NaryExpr, And, Or, Concat
@@ -272,6 +273,7 @@
272273
"ExtractUint16",
273274
"ExtractUint32",
274275
"ExtractUint64",
276+
"Replace",
275277
"Log",
276278
"While",
277279
"For",

‎pyteal/ast/replace.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
from typing import cast, TYPE_CHECKING
2+
3+
from pyteal.types import TealType, require_type
4+
from pyteal.errors import verifyTealVersion
5+
from pyteal.ir import TealOp, Op, TealBlock
6+
from pyteal.ast.expr import Expr
7+
from pyteal.ast.int import Int
8+
from pyteal.ast.ternaryexpr import TernaryExpr
9+
10+
if TYPE_CHECKING:
11+
from pyteal.compiler import CompileOptions
12+
13+
14+
class ReplaceExpr(Expr):
15+
"""An expression for replacing a section of a byte string at a given start index"""
16+
17+
def __init__(self, original: Expr, start: Expr, replacement: Expr) -> None:
18+
super().__init__()
19+
20+
require_type(original, TealType.bytes)
21+
require_type(start, TealType.uint64)
22+
require_type(replacement, TealType.bytes)
23+
24+
self.original = original
25+
self.start = start
26+
self.replacement = replacement
27+
28+
# helper method for correctly populating op
29+
def __get_op(self, options: "CompileOptions"):
30+
s = cast(Int, self.start).value
31+
if s < 2**8:
32+
return Op.replace2
33+
else:
34+
return Op.replace3
35+
36+
def __teal__(self, options: "CompileOptions"):
37+
if not isinstance(self.start, Int):
38+
return TernaryExpr(
39+
Op.replace3,
40+
(TealType.bytes, TealType.uint64, TealType.bytes),
41+
TealType.bytes,
42+
self.original,
43+
self.start,
44+
self.replacement,
45+
).__teal__(options)
46+
47+
op = self.__get_op(options)
48+
49+
verifyTealVersion(
50+
op.min_version,
51+
options.version,
52+
"TEAL version too low to use op {}".format(op),
53+
)
54+
55+
s = cast(Int, self.start).value
56+
if op == Op.replace2:
57+
return TealBlock.FromOp(
58+
options, TealOp(self, op, s), self.original, self.replacement
59+
)
60+
elif op == Op.replace3:
61+
return TealBlock.FromOp(
62+
options, TealOp(self, op), self.original, self.start, self.replacement
63+
)
64+
65+
def __str__(self):
66+
return "(Replace {} {} {})".format(self.original, self.start, self.replacement)
67+
68+
def type_of(self):
69+
return TealType.bytes
70+
71+
def has_return(self):
72+
return False
73+
74+
75+
def Replace(original: Expr, start: Expr, replacement: Expr) -> Expr:
76+
"""
77+
Replace a portion of original bytes with new bytes at a given starting point.
78+
79+
Requires TEAL version 7 or higher.
80+
81+
Args:
82+
original: The value containing the original bytes. Must evaluate to bytes.
83+
start: The index of the byte where replacement starts. Must evaluate to an integer less than Len(original).
84+
replacement: The value containing the replacement bytes. Must evaluate to bytes with length at most Len(original) - start.
85+
"""
86+
return ReplaceExpr(original, start, replacement)

‎pyteal/ast/replace_test.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import pytest
2+
3+
import pyteal as pt
4+
5+
teal6Options = pt.CompileOptions(version=6)
6+
teal7Options = pt.CompileOptions(version=7)
7+
8+
9+
def test_replace_immediate():
10+
args = [pt.Bytes("my string"), pt.Int(0), pt.Bytes("abcdefghi")]
11+
expr = pt.Replace(args[0], args[1], args[2])
12+
assert expr.type_of() == pt.TealType.bytes
13+
14+
expected = pt.TealSimpleBlock(
15+
[
16+
pt.TealOp(args[0], pt.Op.byte, '"my string"'),
17+
pt.TealOp(args[2], pt.Op.byte, '"abcdefghi"'),
18+
pt.TealOp(expr, pt.Op.replace2, 0),
19+
]
20+
)
21+
22+
actual, _ = expr.__teal__(teal7Options)
23+
actual.addIncoming()
24+
actual = pt.TealBlock.NormalizeBlocks(actual)
25+
26+
assert actual == expected
27+
28+
with pytest.raises(pt.TealInputError):
29+
expr.__teal__(teal6Options)
30+
31+
32+
def test_replace_stack_int():
33+
my_string = "*" * 257
34+
args = [pt.Bytes(my_string), pt.Int(256), pt.Bytes("ab")]
35+
expr = pt.Replace(args[0], args[1], args[2])
36+
assert expr.type_of() == pt.TealType.bytes
37+
38+
expected = pt.TealSimpleBlock(
39+
[
40+
pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
41+
pt.TealOp(args[1], pt.Op.int, 256),
42+
pt.TealOp(args[2], pt.Op.byte, '"ab"'),
43+
pt.TealOp(expr, pt.Op.replace3),
44+
]
45+
)
46+
47+
actual, _ = expr.__teal__(teal7Options)
48+
actual.addIncoming()
49+
actual = pt.TealBlock.NormalizeBlocks(actual)
50+
51+
assert actual == expected
52+
53+
with pytest.raises(pt.TealInputError):
54+
expr.__teal__(teal6Options)
55+
56+
57+
# Mirrors `test_replace_stack_int`, but attempts replacement with start != pt.Int.
58+
def test_replace_stack_not_int():
59+
my_string = "*" * 257
60+
add = pt.Add(pt.Int(254), pt.Int(2))
61+
args = [pt.Bytes(my_string), add, pt.Bytes("ab")]
62+
expr = pt.Replace(args[0], args[1], args[2])
63+
assert expr.type_of() == pt.TealType.bytes
64+
65+
expected = pt.TealSimpleBlock(
66+
[
67+
pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
68+
pt.TealOp(pt.Int(254), pt.Op.int, 254),
69+
pt.TealOp(pt.Int(2), pt.Op.int, 2),
70+
pt.TealOp(add, pt.Op.add),
71+
pt.TealOp(args[2], pt.Op.byte, '"ab"'),
72+
pt.TealOp(expr, pt.Op.replace3),
73+
]
74+
)
75+
76+
actual, _ = expr.__teal__(teal7Options)
77+
actual.addIncoming()
78+
actual = pt.TealBlock.NormalizeBlocks(actual)
79+
80+
with pt.TealComponent.Context.ignoreExprEquality():
81+
assert actual == expected
82+
83+
with pytest.raises(pt.TealInputError):
84+
expr.__teal__(teal6Options)
85+
86+
87+
def test_replace_invalid():
88+
with pytest.raises(pt.TealTypeError):
89+
pt.Replace(pt.Bytes("my string"), pt.Int(0), pt.Int(1))
90+
91+
with pytest.raises(pt.TealTypeError):
92+
pt.Replace(
93+
pt.Bytes("my string"), pt.Bytes("should be int"), pt.Bytes("abcdefghi")
94+
)
95+
96+
with pytest.raises(pt.TealTypeError):
97+
pt.Replace(pt.Bytes("my string"), pt.Txn.sender(), pt.Bytes("abcdefghi"))

‎pyteal/ast/substring.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def __init__(self, stringArg: Expr, startArg: Expr, endArg: Expr) -> None:
2626
self.endArg = endArg
2727

2828
# helper method for correctly populating op
29-
def __getOp(self, options: "CompileOptions"):
29+
def __get_op(self, options: "CompileOptions"):
3030
s, e = cast(Int, self.startArg).value, cast(Int, self.endArg).value
3131
l = e - s
3232

@@ -58,7 +58,7 @@ def __teal__(self, options: "CompileOptions"):
5858
self.endArg,
5959
).__teal__(options)
6060

61-
op = self.__getOp(options)
61+
op = self.__get_op(options)
6262

6363
verifyTealVersion(
6464
op.min_version,
@@ -121,7 +121,7 @@ def __init__(self, stringArg: Expr, startArg: Expr, lenArg: Expr) -> None:
121121
self.lenArg = lenArg
122122

123123
# helper method for correctly populating op
124-
def __getOp(self, options: "CompileOptions"):
124+
def __get_op(self, options: "CompileOptions"):
125125
s, l = cast(Int, self.startArg).value, cast(Int, self.lenArg).value
126126
if s < 2**8 and l > 0 and l < 2**8:
127127
return Op.extract
@@ -139,7 +139,7 @@ def __teal__(self, options: "CompileOptions"):
139139
self.lenArg,
140140
).__teal__(options)
141141

142-
op = self.__getOp(options)
142+
op = self.__get_op(options)
143143

144144
verifyTealVersion(
145145
op.min_version,
@@ -186,7 +186,7 @@ def __init__(
186186
self.startArg = startArg
187187

188188
# helper method for correctly populating op
189-
def __getOp(self, options: "CompileOptions"):
189+
def __get_op(self, options: "CompileOptions"):
190190
if not isinstance(self.startArg, Int):
191191
return Op.substring3
192192

@@ -197,7 +197,7 @@ def __getOp(self, options: "CompileOptions"):
197197
return Op.substring3
198198

199199
def __teal__(self, options: "CompileOptions"):
200-
op = self.__getOp(options)
200+
op = self.__get_op(options)
201201

202202
verifyTealVersion(
203203
op.min_version,

‎pyteal/ast/substring_test.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,88 @@ def test_suffix_stack():
272272
assert actual == expected
273273

274274

275+
@pytest.mark.parametrize("op", [pt.Op.extract3, pt.Op.substring3])
276+
def test_startArg_not_int(op: pt.Op):
277+
my_string = "*" * 257
278+
add = pt.Add(pt.Int(254), pt.Int(2))
279+
args = [pt.Bytes(my_string), add, pt.Int(257)]
280+
281+
def generate_expr() -> pt.Expr:
282+
match op:
283+
case pt.Op.extract3:
284+
return pt.Extract(args[0], args[1], args[2])
285+
case pt.Op.substring3:
286+
return pt.Substring(args[0], args[1], args[2])
287+
case _:
288+
raise Exception(f"Unsupported {op=}")
289+
290+
expr = generate_expr()
291+
assert expr.type_of() == pt.TealType.bytes
292+
293+
expected = pt.TealSimpleBlock(
294+
[
295+
pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
296+
pt.TealOp(pt.Int(254), pt.Op.int, 254),
297+
pt.TealOp(pt.Int(2), pt.Op.int, 2),
298+
pt.TealOp(add, pt.Op.add),
299+
pt.TealOp(args[2], pt.Op.int, 257),
300+
pt.TealOp(None, op),
301+
]
302+
)
303+
304+
actual, _ = expr.__teal__(teal5Options)
305+
actual.addIncoming()
306+
actual = pt.TealBlock.NormalizeBlocks(actual)
307+
308+
with pt.TealComponent.Context.ignoreExprEquality():
309+
assert actual == expected
310+
311+
if op == pt.Op.extract3:
312+
with pytest.raises(pt.TealInputError):
313+
expr.__teal__(teal4Options)
314+
315+
316+
@pytest.mark.parametrize("op", [pt.Op.extract3, pt.Op.substring3])
317+
def test_endArg_not_int(op: pt.Op):
318+
my_string = "*" * 257
319+
add = pt.Add(pt.Int(254), pt.Int(3))
320+
args = [pt.Bytes(my_string), pt.Int(256), add]
321+
322+
def generate_expr() -> pt.Expr:
323+
match op:
324+
case pt.Op.extract3:
325+
return pt.Extract(args[0], args[1], args[2])
326+
case pt.Op.substring3:
327+
return pt.Substring(args[0], args[1], args[2])
328+
case _:
329+
raise Exception(f"Unsupported {op=}")
330+
331+
expr = generate_expr()
332+
assert expr.type_of() == pt.TealType.bytes
333+
334+
expected = pt.TealSimpleBlock(
335+
[
336+
pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
337+
pt.TealOp(args[1], pt.Op.int, 256),
338+
pt.TealOp(pt.Int(254), pt.Op.int, 254),
339+
pt.TealOp(pt.Int(3), pt.Op.int, 3),
340+
pt.TealOp(add, pt.Op.add),
341+
pt.TealOp(None, op),
342+
]
343+
)
344+
345+
actual, _ = expr.__teal__(teal5Options)
346+
actual.addIncoming()
347+
actual = pt.TealBlock.NormalizeBlocks(actual)
348+
349+
with pt.TealComponent.Context.ignoreExprEquality():
350+
assert actual == expected
351+
352+
if op == pt.Op.extract3:
353+
with pytest.raises(pt.TealInputError):
354+
expr.__teal__(teal4Options)
355+
356+
275357
def test_suffix_invalid():
276358
with pytest.raises(pt.TealTypeError):
277359
pt.Suffix(pt.Int(0), pt.Int(0))

‎pyteal/ir/ops.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,8 @@ def min_version(self) -> int:
179179
gitxnas = OpType("gitxnas", Mode.Application, 6)
180180
gloadss = OpType("gloadss", Mode.Application, 6)
181181
acct_params_get = OpType("acct_params_get", Mode.Application, 6)
182+
replace2 = OpType("replace2", Mode.Signature | Mode.Application, 7)
183+
replace3 = OpType("replace3", Mode.Signature | Mode.Application, 7)
182184
# fmt: on
183185

184186

0 commit comments

Comments
 (0)