remove wrap_typ util and replace w if-else chain

tests/functional/builtins/codegen/test_extract32.py

@@ -1,6 +1,5 @@
 import pytest
 
-from tests.utils import wrap_typ_with_storage_loc
 from vyper.evm.opcodes import version_check
 
 
@@ -10,8 +9,12 @@ def test_extract32_extraction(tx_failed, get_contract_with_gas_estimation, locat
         pytest.skip(
             "Skipping test as storage_location is 'transient' and EVM version is pre-Cancun"
         )
+    if location == "storage":
+        decl = "y: Bytes[100]"
+    else:
+        decl = "y: transient(Bytes[100])"
     extract32_code = f"""
-y: {wrap_typ_with_storage_loc("Bytes[100]", location)}
+{decl}
 @external
 def extrakt32(inp: Bytes[100], index: uint256) -> bytes32:
     return extract32(inp, index)

tests/functional/builtins/codegen/test_slice.py

@@ -2,7 +2,6 @@
 import pytest
 from hypothesis import given, settings
 
-from tests.utils import wrap_typ_with_storage_loc
 from vyper.compiler import compile_code
 from vyper.compiler.settings import OptimizationLevel, Settings
 from vyper.evm.opcodes import version_check
@@ -124,15 +123,15 @@ def test_slice_bytes_fuzz(
     if location == "memory":
         spliced_code = f"foo: Bytes[{length_bound}] = inp"
         foo = "foo"
-    elif location == "transient" or location == "storage":
+    elif location == "storage":
         preamble = f"""
-foo: {wrap_typ_with_storage_loc(f"Bytes[{length_bound}]", location)}
+foo: Bytes[{length_bound}]
          """
         spliced_code = "self.foo = inp"
         foo = "self.foo"
-    elif location == "storage":
+    elif location == "transient":
         preamble = f"""
-foo: Bytes[{length_bound}]
+foo: transient(Bytes[{length_bound}])
         """
         spliced_code = "self.foo = inp"
         foo = "self.foo"
@@ -209,12 +208,22 @@ def _get_contract():
 
 
 @pytest.mark.parametrize(
-    "location", ["storage"] if not version_check(begin="cancun") else ["storage", "transient"]
+    "location", ["storage", "transient"]
 )
 def test_slice_private(get_contract, location):
+    if location == "transient" and not version_check(begin="cancun"):
+        pytest.skip(
+            "Skipping test as storage_location is 'transient' and EVM version is pre-Cancun"
+        )
+
     # test there are no buffer overruns in the slice function
+    if location == "storage":
+        decl = "bytez: public(String[12])"
+    else:
+        decl = "bytez: public(transient(String[12]))"
+
     code = f"""
-bytez: public({wrap_typ_with_storage_loc("String[12]", location)})
+{decl}
 
 @internal
 def _slice(start: uint256, length: uint256):

tests/utils.py

@@ -19,11 +19,3 @@ def parse_and_fold(source_code):
     ast = vy_ast.parse_to_ast(source_code)
     constant_fold(ast)
     return ast
-
-
-def wrap_typ_with_storage_loc(typ, loc):
-    if loc == "storage":
-        return typ
-    elif loc == "transient":
-        return f"transient({typ})"
-    raise AssertionError(f"unreachable storage location {loc}")  # pragma: nocover