Refactor to update methods

Extract common behaviour in different functions

wordle_env/state.py

@@ -64,12 +64,7 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
         offset = 1 + cint * WORDLE_N * 3
         if mask[i] == YES:
             prior_yes.append(c)
-            # char at position i = yes, all other chars at position i == no
-            state[offset + 3 * i:offset + 3 * i + 3] = [0, 0, 1]
-            for ocint in range(len(WORDLE_CHARS)):
-                if ocint != cint:
-                    oc_offset = 1 + ocint * WORDLE_N * 3
-                    state[oc_offset + 3 * i:oc_offset + 3 * i + 3] = [1, 0, 0]
+            _set_yes(state, offset, cint, i)
 
     for i, c in enumerate(word):
         cint = ord(c) - ord(WORDLE_CHARS[0])
@@ -77,11 +72,8 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
         if mask[i] == SOMEWHERE:
             prior_maybe.append(c)
             # Char at position i = no, and in other positions maybe except it had a value before, other chars stay as they are
-            for char_idx in range(0, WORDLE_N * 3, 3):
-                char_offset = offset + char_idx
-                if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
-                    state[char_offset: char_offset + 3] = [0, 1, 0]
-            state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
+            _set_no(state, offset, i)
+            _set_if_cero(state, offset, [0, 1, 0])
         elif mask[i] == NO:
             # Need to check this first in case there's prior maybe + yes
             if c in prior_maybe:
@@ -95,7 +87,7 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
                         state[offset + 3 * j:offset + 3 * j + 3] = [1, 0, 0]
             else:
                 # Just straight up no
-                state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
+                _set_all_no(state, offset)
     return state
 
 
@@ -148,11 +140,7 @@ def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState,
         if goal_word[i] == c:
             # char at position i = yes, all other chars at position i == no
             reward += CHAR_REWARD
-            state[offset + 3 * i:offset + 3 * i + 3] = [0, 0, 1]
-            for ocint in range(len(WORDLE_CHARS)):
-                if ocint != cint:
-                    oc_offset = 1 + ocint * WORDLE_N * 3
-                    state[oc_offset + 3 * i:oc_offset + 3 * i + 3] = [1, 0, 0]
+            _set_yes(state, offset, cint, i)
             processed_letters.append(c)
 
     for i, c in enumerate(word):
@@ -161,21 +149,40 @@ def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState,
         if goal_word[i] != c:
             if c in goal_word and goal_word.count(c) > processed_letters.count(c):
                 # Char at position i = no, and in other positions maybe except it had a value before, other chars stay as they are
-                for char_idx in range(0, WORDLE_N * 3, 3):
-                    char_offset = offset + char_idx
-                    if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
-                        state[char_offset: char_offset + 3] = [0, 1, 0]
-                state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
+                _set_no(state, offset, i)
+                _set_if_cero(state, offset, [0, 1, 0])
                 reward += CHAR_REWARD * 0.1
             elif c not in goal_word:
                 # Char at all positions = no
-                state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
+                _set_all_no(state, offset)
             else: # goal_word.count(c) <= processed_letters.count(c) and goal in word
                 # At i and in every position which is not set = no
-                state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
-                for char_idx in range(0, WORDLE_N * 3, 3):
-                    char_offset = offset + char_idx
-                    if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
-                        state[char_offset: char_offset + 3] = [1, 0, 0]
+                _set_no(state, offset, i)
+                _set_if_cero(state, offset, [1, 0, 0])
             processed_letters.append(c)
     return state, reward
+
+
+def _set_if_cero(state, offset, value):
+    for char_idx in range(0, WORDLE_N * 3, 3):
+        char_offset = offset + char_idx
+        if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
+            state[char_offset: char_offset + 3] = value
+
+
+def _set_yes(state, offset, char_int, char_pos):
+    # char at position char_pos = yes, all other chars at position char_pos == no
+    pos_offset = 3 * char_pos
+    state[offset + pos_offset:offset + pos_offset + 3] = [0, 0, 1]
+    for ocint in range(len(WORDLE_CHARS)):
+        if ocint != char_int:
+            oc_offset = 1 + ocint * WORDLE_N * 3
+            state[oc_offset + pos_offset:oc_offset + pos_offset + 3] = [1, 0, 0]
+
+
+def _set_no(state, offset, char_pos):
+    state[offset + 3 * char_pos:offset + 3 * char_pos + 3] = [1, 0, 0]
+
+
+def _set_all_no(state, offset):
+    state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N