Data recording help for key:value-like states.

docs/source/user_guide/tutorials/data.rst

@@ -128,8 +128,11 @@ Note that the string State of ``people_seen`` acts as a way to record data, even
 the moment the name of the last scanned person. This lets states behave as carriers of current or past
 information, depending on your needs.
 
-The ``items`` value doesn't record, because the state doesn't see ``cash.items = ...``.
-For objects like that, you should:
+Complex States
+--------------
+
+The ``items`` value doesn't record, because the state doesn't see the ``cash.items = ...`` operation.
+For objects like that, you can use the ``record_state`` method on the ``Actor``:
 
 .. code:: python
 
@@ -141,22 +144,23 @@ For objects like that, you should:
     with UP.EnvironmentContext() as env:
         cash = Cashier(name="Ertha")
         cash.items["bread"] = 1
-        cash.items = cash.items # <- Tell the state it's been changed explicitly
+        cash.record_state("items")
+        # or, cash.items = cash.items
         env.run(until=0.75)
         cash.items["bread"] += 2
         cash.items["milk"] += 3
-        cash.items = cash.items
+        cash.record_state("items")
 
         print(cash._state_histories)
         >>>{'items': [(0.0, Counter({'bread': 1})), (0.75, Counter({'bread': 3, 'milk': 3}))]}
 
-This is clunky, but the ``Counter`` object has no way of knowing it belongs in a ``State`` to get the
-recording to work. In the future, UPSTAGE may monkey-patch objects with ``__set__`` methods, but for
-now this is the workaround.
-
 Note also that UPSTAGE deep-copies the value in the state history, so any data should be compatible with that
 operation.
 
+UPSTAGE will output data from ``dataclass`` states, and ``dict[str, Any]`` states by creating rows in the
+data table with the naming convention ``state_name.attribute_name``, where the attribute is either a dataclass
+attribute or a key from the dictionary.
+
 Geographic Types
 ----------------
 

src/upstage_des/actor.py

@@ -1148,5 +1148,19 @@ def get_nucleus(self) -> "TaskNetworkNucleus":
             raise SimulationError("Expected a nucleus, but none found.")
         return self._state_listener
 
+    def record_state(self, state_name: str) -> None:
+        """Record a state by its name.
+
+        Useful for states that have attributes that aren't set
+        via the descriptor, such as dictionaries or dataclasses.
+
+        Args:
+            state_name (str): The name of the state.
+        """
+        if state_name not in self.states:
+            raise SimulationError(f"No state '{state_name}' to record.")
+        v = getattr(self, state_name)
+        self._state_defs[state_name]._do_record(self, v)
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}: {self.name}"

src/upstage_des/data_utils.py

@@ -1,5 +1,6 @@
 """Utilities for gathering all recorded simulation data."""
 
+from dataclasses import asdict, is_dataclass
 from typing import Any, cast
 
 from upstage_des.actor import Actor
@@ -49,11 +50,21 @@ def _state_history_to_table(
     for time, value in hist:
         if isinstance(value, ActiveStatus):
             row = data.pop(-1)
-            row = tuple(list(row[:-1]) + [value.name])
+            rows = [tuple(list(row[:-1]) + [value.name])]
             active_value = "active" if value.name == "activating" else "inactive"
+        elif is_dataclass(value) and not isinstance(value, type):
+            rows = [
+                (actor_name, actor_kind, f"{state_name}.{k}", time, v, active_value)
+                for k, v in asdict(value).items()
+            ]
+        elif isinstance(value, dict):
+            rows = [
+                (actor_name, actor_kind, f"{state_name}.{k}", time, v, active_value)
+                for k, v in value.items()
+            ]
         else:
-            row = (actor_name, actor_kind, state_name, time, value, active_value)
-        data.append(row)
+            rows = [(actor_name, actor_kind, state_name, time, value, active_value)]
+        data.extend(rows)
     return data
 
 

src/upstage_des/test/test_data_reporting.py

@@ -1,23 +1,32 @@
 """Test the data recording/reporting capabilities."""
 
 from collections import Counter
+from dataclasses import dataclass
 
 import upstage_des.api as UP
 from upstage_des.data_utils import create_location_table, create_table
 
 
+@dataclass
+class Information:
+    value_1: int
+    value_2: float
+
+
 class Cashier(UP.Actor):
     items_scanned = UP.State[int](recording=True)
     other = UP.State[float]()
     cue = UP.State[UP.SelfMonitoringStore]()
     cue2 = UP.ResourceState[UP.SelfMonitoringContainer](default=UP.SelfMonitoringContainer)
     time_working = UP.LinearChangingState(default=0.0, recording=True, record_duplicates=True)
+    info = UP.State[Information](recording=True)
 
 
 class Cart(UP.Actor):
     location = UP.CartesianLocationChangingState(recording=True)
     location_two = UP.CartesianLocationChangingState(recording=True)
     holding = UP.State[float](default=0.0, recording=True)
+    some_data = UP.State[dict[str, float]](recording=True)
 
 
 def test_data_reporting() -> None:
@@ -28,19 +37,22 @@ def test_data_reporting() -> None:
             other=0.0,
             items_scanned=0,
             cue=UP.SelfMonitoringStore(env),
+            info=Information(0, 0),
         )
 
         cash2 = Cashier(
             name="Bertha",
             other=0.0,
             items_scanned=0,
             cue=UP.SelfMonitoringStore(env),
+            info=Information(1, 1),
         )
         store = UP.SelfMonitoringFilterStore(env, name="Store Test")
         cart = Cart(
             name="Wobbly Wheel",
             location=UP.CartesianLocation(1.0, 1.0),
             location_two=UP.CartesianLocation(1.0, 1.0),
+            some_data={"exam": 2.0},
         )
 
         for c in [cash, cash2]:
@@ -76,6 +88,10 @@ def test_data_reporting() -> None:
         cart.location_two
         cash.items_scanned += 2
         store.put("XYZ")
+        cash.info.value_1 = 3
+        cash.record_state("info")
+        cash2.info.value_2 = 4.3
+        cash2.record_state("info")
 
         for c in [cash, cash2]:
             c.cue.put("B")
@@ -88,6 +104,8 @@ def test_data_reporting() -> None:
         env.run(until=3)
         cart.location
         cart.location_two
+        cart.some_data["exam"] = 4.0
+        cart.record_state("some_data")
 
         cart.deactivate_state(state="location", task=t)
         cart.deactivate_state(state="location_two", task=t)
@@ -103,6 +121,8 @@ def test_data_reporting() -> None:
         env.run(until=3.3)
         cart.location
         cart.location_two = UP.CartesianLocation(-1.0, -1.0)
+        cart.some_data["new"] = 123.45
+        cart.record_state("some_data")
         store.put("ABC")
         env.run()
         cart.location
@@ -124,13 +144,19 @@ def test_data_reporting() -> None:
     assert ctr[("Bertha", "Cashier", "cue2")] == 4
     assert ctr[("Bertha", "Cashier", "time_working")] == 5
     assert ctr[("Store Test", "SelfMonitoringFilterStore", "Resource")] == 3
+    assert ctr[("Ertha", "Cashier", "info.value_1")] == 2
+    assert ctr[("Ertha", "Cashier", "info.value_2")] == 2
+    assert ctr[("Bertha", "Cashier", "info.value_1")] == 2
+    assert ctr[("Bertha", "Cashier", "info.value_2")] == 2
     # Test for default values untouched in the sim showing up in the data.
     assert ctr[("Wobbly Wheel", "Cart", "holding")] == 1
+    assert ctr[("Wobbly Wheel", "Cart", "some_data.exam")] == 3
+    assert ctr[("Wobbly Wheel", "Cart", "some_data.new")] == 1
     row = [r for r in state_table if r[:3] == ("Wobbly Wheel", "Cart", "holding")][0]
     assert row[4] == 0
     assert row[3] == 0.0
     # Continuing as before
-    assert len(state_table) == 38
+    assert len(state_table) == 50
     assert cols == all_cols
     assert cols == [
         "Entity Name",
@@ -154,7 +180,7 @@ def test_data_reporting() -> None:
     assert ctr[("Wobbly Wheel", "Cart", "holding")] == 1
     assert ctr[("Wobbly Wheel", "Cart", "location")] == 4
     assert ctr[("Wobbly Wheel", "Cart", "location_two")] == 4
-    assert len(all_state_table) == 38 + 8
+    assert len(all_state_table) == 38 + 8 + 12
 
     assert loc_cols == [
         "Entity Name",