From 94a196d0dbe79f4ae205d13f50302316860194b8 Mon Sep 17 00:00:00 2001
From: jk-ethz <80681863+jk-ethz@users.noreply.github.com>
Date: Wed, 16 Oct 2024 18:29:34 +0200
Subject: [PATCH] Add userdata constants defined outside state code (#90)

* Add userdata constants

* Revert a change potentially made by mistake

---------

Co-authored-by: Isaac Saito <130s@users.noreply.github.com>
---
 smach/src/smach/state_machine.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/smach/src/smach/state_machine.py b/smach/src/smach/state_machine.py
index dcfb61f..5376cf5 100644
--- a/smach/src/smach/state_machine.py
+++ b/smach/src/smach/state_machine.py
@@ -62,6 +62,7 @@ def __init__(self, outcomes, input_keys=[], output_keys=[]):
         self._states = {}
         self._transitions = {}
         self._remappings = {}
+        self._constants = {}
 
         # Construction vars
         self._last_added_label = None
@@ -81,7 +82,7 @@ def __setstate__(self, d):
 
     ### Construction methods
     @staticmethod
-    def add(label, state, transitions=None, remapping=None):
+    def add(label, state, transitions=None, remapping=None, constants=None):
         """Add a state to the opened state machine.
         
         @type label: string
@@ -92,8 +93,11 @@ def add(label, state, transitions=None, remapping=None):
         @param transitions: A dictionary mapping state outcomes to other state
         labels or container outcomes.
 
-        @param remapping: A dictrionary mapping local userdata keys to userdata
+        @param remapping: A dictionary mapping local userdata keys to userdata
         keys in the container.
+
+        @param constants: A dictionary mapping userdata keys in the container
+        to constant values.
         """
         # Get currently opened container
         self = StateMachine._currently_opened_container()
@@ -110,6 +114,9 @@ def add(label, state, transitions=None, remapping=None):
         if remapping is None:
             remapping = {}
 
+        if constants is None:
+            constants = {}
+
         # Add group transitions to this new state, if they exist
         """
         if 'transitions' in smach.Container._context_kwargs:
@@ -141,6 +148,7 @@ def add(label, state, transitions=None, remapping=None):
         self._states[label] = state
         self._transitions[label] = transitions
         self._remappings[label] = remapping
+        self._constants[label] = constants
         smach.logdebug("TRANSITIONS FOR %s: %s" % (label, str(self._transitions[label])))
 
         # Add transition to this state if connected outcome is defined
@@ -154,7 +162,7 @@ def add(label, state, transitions=None, remapping=None):
         return state
 
     @staticmethod
-    def add_auto(label, state, connector_outcomes, transitions=None, remapping=None):
+    def add_auto(label, state, connector_outcomes, transitions=None, remapping=None, constants=None):
         """Add a state to the state machine such that it automatically
         transitions to the next added state.
 
@@ -180,7 +188,7 @@ def add_auto(label, state, connector_outcomes, transitions=None, remapping=None)
         self = StateMachine._currently_opened_container()
 
         # First add this state
-        add_ret = smach.StateMachine.add(label, state, transitions, remapping)
+        add_ret = smach.StateMachine.add(label, state, transitions, remapping, constants)
 
         # Make sure the connector outcomes are valid for this state
         registered_outcomes = state.get_registered_outcomes()
@@ -247,6 +255,8 @@ def _update_once(self):
         # Execute the state
         try:
             self._state_transitioning_lock.release()
+            for k, v in self._constants[self._current_label].items():
+                self.userdata[k] = v
             outcome = self._current_state.execute(
                     smach.Remapper(
                         self.userdata,