From d354d4e14a88b68f25782a74e48bd8d30d81690e Mon Sep 17 00:00:00 2001
From: Daniel McGregor <daniel.mcgregor@makerx.com.au>
Date: Mon, 8 Apr 2024 17:47:15 +0800
Subject: [PATCH] fix: prevent usage of state proxies (GlobalState, LocalState)
 outside __init__ method

BREAKING CHANGE: Using state proxies (GlobalState, LocalState) outside of an __init__ method may not give the behaviour expected, so prevent their usage in those scenarios
---
 src/puya/awst_build/constants.py              |  4 ++--
 src/puya/awst_build/subroutine.py             |  5 +++++
 .../test_expected_output/expected_errors.test | 21 +++++++++++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/puya/awst_build/constants.py b/src/puya/awst_build/constants.py
index 435dd3e09d..b601397a13 100644
--- a/src/puya/awst_build/constants.py
+++ b/src/puya/awst_build/constants.py
@@ -29,9 +29,9 @@
 CLS_BIGUINT_ALIAS = f"{ALGOPY_PREFIX}BigUInt"
 CLS_TRANSACTION_BASE = f"{ALGOPY_PREFIX}gtxn.TransactionBase"
 CLS_LOCAL_STATE = f"{ALGOPY_PREFIX}_state.LocalState"
-CLS_LOCAL_STATE_ALIAS = f"{ALGOPY_PREFIX}.LocalState"
+CLS_LOCAL_STATE_ALIAS = f"{ALGOPY_PREFIX}LocalState"
 CLS_GLOBAL_STATE = f"{ALGOPY_PREFIX}_state.GlobalState"
-CLS_GLOBAL_STATE_ALIAS = f"{ALGOPY_PREFIX}.GlobalState"
+CLS_GLOBAL_STATE_ALIAS = f"{ALGOPY_PREFIX}GlobalState"
 SUBROUTINE_HINT = f"{ALGOPY_PREFIX}_hints.subroutine"
 LOGICSIG_DECORATOR = f"{ALGOPY_PREFIX}_logic_sig.logicsig"
 LOGICSIG_DECORATOR_ALIAS = f"{ALGOPY_PREFIX}.logicsig"
diff --git a/src/puya/awst_build/subroutine.py b/src/puya/awst_build/subroutine.py
index 9ae828830b..9e1aac9021 100644
--- a/src/puya/awst_build/subroutine.py
+++ b/src/puya/awst_build/subroutine.py
@@ -347,6 +347,11 @@ def _handle_state_proxy_assignment(
             raise CodeError(
                 f"{rvalue.python_name} should only be used inside a contract class", stmt_loc
             )
+        if self.func_def.name != "__init__":
+            raise CodeError(
+                f"{rvalue.python_name} can only be used in the __init__ method",
+                stmt_loc,
+            )
         if len(lvalues) != 1:
             raise CodeError(
                 f"{rvalue.python_name} can only be assigned to a single member variable",
diff --git a/tests/test_expected_output/expected_errors.test b/tests/test_expected_output/expected_errors.test
index e2f9d02695..3df5f9ef2c 100644
--- a/tests/test_expected_output/expected_errors.test
+++ b/tests/test_expected_output/expected_errors.test
@@ -135,3 +135,24 @@ class Baddie(arc4.ARC4Contract):
         arr2.append(arc4.Byte(1))
         assert foo
 
+
+## case: test_state_proxies_outside_init_fail
+from algopy import GlobalState, LocalState, Txn, UInt64, arc4
+
+class Baddie(arc4.ARC4Contract):
+    def __init__(self) -> None:
+        self.global_a = GlobalState(UInt64)
+        self.local_b = LocalState(UInt64)
+
+    @arc4.abimethod
+    def okay(self) -> None:
+        self.global_a.value = UInt64(123)
+        self.local_b[Txn.sender] = UInt64(456)
+
+    @arc4.abimethod()
+    def not_okay1(self) -> None:
+        self.global_c = GlobalState(UInt64) ## E: algopy.GlobalState can only be used in the __init__ method
+
+    @arc4.abimethod()
+    def not_okay2(self) -> None:
+        self.local_d = LocalState(UInt64) ## E: algopy.LocalState can only be used in the __init__ method
\ No newline at end of file