Merge pull request #278 from Point72/tkp/staterework

rework state to remove annoyances and inconsistensies, prep for Stage and UI APIs

Author: timkpaine

.github/workflows/build.yaml

@@ -111,13 +111,13 @@ jobs:
       uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
       with:
         name: test-results-${{ matrix.os }}-${{ matrix.python-version }}
-        path: '**/junit.xml'
+        path: junit.xml
       if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
 
     - name: Publish Unit Test Results
       uses: EnricoMi/publish-unit-test-result-action@c950f6fb443cb5af20a377fd0dfaa78838901040 # v2
       with:
-        files: '**/junit.xml'
+        files: junit.xml
       if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
 
     - name: Upload coverage

Makefile

@@ -97,7 +97,7 @@ test-py:  ## run python tests
 tests-py: test-py
 
 coverage-py:  ## run python tests and collect test coverage
-	python -m pytest -v csp_gateway/tests --cov=csp_gateway --cov-report term-missing --cov-report xml
+	python -m pytest -v csp_gateway/tests --junitxml=junit.xml --cov=csp_gateway --cov-report term-missing --cov-report xml
 
 .PHONY: test-js tests-js coverage-js
 test-js:  ## run js tests

csp_gateway/server/demo/omnibus.py

@@ -115,7 +115,14 @@ class ExampleGatewayChannels(GatewayChannels):
     example_list: ts[List[ExampleData]] = None
     never_ticks: ts[ExampleData] = None
 
-    s_example: ts[State[ExampleData]] = None
+    # State fields can be added via annotation or the `set_state` API in the module's `connect` method
+    example_with_state: Annotated[ts[ExampleData], State(("id", "x"))] = None
+    example_with_state_multiple: Annotated[
+        ts[ExampleData],
+        State(("id", "x")),
+        State(("id", "y"), alias="example_with_state_alternative"),
+    ] = None
+    # example: Set in connect with a different schema, to show flexibility of state definition
 
     basket: Dict[ExampleEnum, ts[ExampleData]] = None
     str_basket: Dict[str, ts[ExampleData]] = None
@@ -179,6 +186,8 @@ def connect(self, channels: ExampleGatewayChannels):
         # Channels set via `set_channel`
         channels.set_channel(ExampleGatewayChannels.example, data)
         channels.set_channel(ExampleGatewayChannels.example_list, data_list)
+        channels.set_channel(ExampleGatewayChannels.example_with_state, data)
+        channels.set_channel(ExampleGatewayChannels.example_with_state_multiple, data)
 
         # Generic channel for sending data from non-csp sources
         channels.add_send_channel(ExampleGatewayChannels.example)
@@ -187,8 +196,11 @@ def connect(self, channels: ExampleGatewayChannels):
         channels.add_send_channel(ExampleGatewayChannels.basket, ExampleEnum.C)
         channels.add_send_channel(ExampleGatewayChannels.basket)
 
-        # Rudimentary state accumulation via `set_state`
-        channels.set_state(ExampleGatewayChannels.example, "id")
+        # State accumulation via `set_state`
+        channels.set_state(
+            ExampleGatewayChannels.example,
+            ("id",),
+        )
 
         # Create some data streams for dict baskets
         data_a = self.subscribe(

csp_gateway/server/gateway/csp/channels.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Generic, List, Optional, Type
+from typing import Generic, List, Optional, Type, get_args, get_origin
 
 from ccflow import BaseModel
 from csp.impl.enum import Enum
@@ -57,6 +57,24 @@ def build(self, channels: ChannelsType) -> ChannelsType:
         if self.block_set_channels_until is not None:
             channels._block_set_channels_until = self.block_set_channels_until
 
+        # Pre-declare dynamically-created state channels so that modules calling get_state
+        # during connect() are independent of the order in which set_state is called by the
+        # owning module. The owning module includes the name in both dynamic_channels()
+        # (which provides the element type) and dynamic_state_channels().
+        for node in enabled_modules:
+            declared = node.dynamic_state_channels() if hasattr(node, "dynamic_state_channels") else None
+            if not declared:
+                continue
+            owned_channels = node.dynamic_channels() or {}
+            for name in declared:
+                if name not in owned_channels:
+                    raise ValueError(
+                        f"Module {type(node).__name__} declared '{name}' in dynamic_state_channels() but did not include it in dynamic_channels()"
+                    )
+                t = owned_channels[name]
+                element_type = get_args(t)[0] if get_origin(t) is list else t
+                channels._declare_dynamic_state(name, element_type)
+
         # Wire in each edge Provider.
         # The implementation of set_channel will handle multiplexing streams
         for node in enabled_modules:

csp_gateway/server/gateway/csp/factory.py

@@ -49,8 +49,13 @@ def dynamic_channels(self) -> Optional[Dict[str, Union[Type[GatewayStruct], Type
         ...
 
     def dynamic_state_channels(self) -> Optional[Set[str]]:
-        """
-        The set of dynamic channels that have state.
+        """The subset of :meth:`dynamic_channels` for which this module will also call
+        :meth:`GatewayChannels.set_state` from within :meth:`connect`.
+
+        Declaring a name here lets *other* modules call :meth:`GatewayChannels.get_state`
+        on the channel from their own ``connect`` regardless of the order in which modules
+        are connected; the returned state edge is bound to the real state node once the
+        owning module's ``connect`` calls ``set_state``.
         """
         ...
 

csp_gateway/server/gateway/csp/module.py

@@ -7,7 +7,7 @@
 from enum import Enum as PyEnum
 from functools import lru_cache
 from pprint import pformat
-from typing import Any, Deque, Dict, List, Tuple, Union
+from typing import Any, Deque, Dict, List, Optional, Tuple, Union
 
 import csp
 import duckdb
@@ -75,6 +75,16 @@ def disable_duckdb_state() -> None:
     _USE_DUCKDB_STATE = False
 
 
+def _resolve_keyby_attr(record: Any, path: str) -> Any:
+    """Resolve a (possibly dotted) attribute path on ``record``, returning ``None`` if any segment is missing."""
+    obj = record
+    for segment in path.split("."):
+        if obj is None:
+            return None
+        obj = getattr(obj, segment, None)
+    return obj
+
+
 class StateType(CoreEnum):
     UNKNOWN = 0
     DEFAULT = 1
@@ -149,7 +159,7 @@ def insert(self, record: Any) -> None:
 
         for subkey in self._keyby:
             # extract the key from the record
-            subkey_to_use = getattr(record, subkey, None)
+            subkey_to_use = _resolve_keyby_attr(record, subkey)
 
             if subkey == self._keyby[-1]:
                 # Put the element there if last
@@ -408,7 +418,7 @@ def insert(self, record: Struct) -> None:
             obj_id = None
             for subkey in self._keyby:
                 # extract the key from the record
-                subkey_to_use = getattr(record, subkey, None)
+                subkey_to_use = _resolve_keyby_attr(record, subkey)
 
                 if subkey == self._keyby[-1]:
                     if subkey_to_use not in place.keys():
@@ -527,23 +537,45 @@ def get_duckdb_schema_struct(cls: Struct) -> Tuple[Dict, bool]:
     return (new_type_info, use_duckdb)
 
 
-# NOTE: NEVER access State object directly, always access through the __class_getitem__ API
+# NOTE: For runtime state instances, always use the State[<typ>] API.
+# State() called directly (with no parameterized typ) is the annotation
+# marker form used in `Annotated[ts[X], State(keyby=..., indexer=..., alias=...)]`.
 class State(BaseState):
-    def __init__(self, keyby: Union[Tuple[str, ...], str] = ("id",)) -> None:
-        """Switch case between different state specializations based on the type of the records"""
+    # Annotation metadata. Set on every instance; consumed by ChannelsMetaclass
+    # when this State is found in a field's Annotated metadata.
+    _meta_keyby: Union[Tuple[str, ...], str] = ("id",)
+    _meta_indexer: Optional[Union[str, int]] = None
+    _meta_alias: Optional[str] = None
+
+    def __init__(
+        self,
+        keyby: Union[Tuple[str, ...], str] = ("id",),
+        indexer: Optional[Union[str, int]] = None,
+        alias: Optional[str] = None,
+    ) -> None:
+        """Initialize a State.
+
+        When ``self._typ`` is set (via ``State[T](...)``), this constructs a runtime
+        state collection and dispatches to the appropriate backend. Otherwise the
+        instance is treated as an annotation marker and only retains its metadata.
+        """
+        # Always retain annotation metadata.
+        self._meta_keyby = keyby
+        self._meta_indexer = indexer
+        self._meta_alias = alias
+
+        typ = getattr(self, "_typ", None)
+        if typ is None:
+            # Annotation-marker form; no runtime backend needed.
+            return
 
         global _USE_DUCKDB_STATE
-        try:
-            typ = self._typ
-            if _USE_DUCKDB_STATE and isinstance(typ, type) and issubclass(typ, Struct):
-                schema, use_duckdb = get_duckdb_schema_struct(typ)
-                if use_duckdb:
-                    self._state_impl = DuckDBState(typ, keyby, schema)
-                    self._state_type = StateType.DUCKDB
-                    return
-        except AttributeError:
-            log.warning("Do not create object directly from State use the State[<typ>] API instead for performance reasons")
-            log.warning("Using DefaultStateClass")
+        if _USE_DUCKDB_STATE and isinstance(typ, type) and issubclass(typ, Struct):
+            schema, use_duckdb = get_duckdb_schema_struct(typ)
+            if use_duckdb:
+                self._state_impl = DuckDBState(typ, keyby, schema)
+                self._state_type = StateType.DUCKDB
+                return
         self._state_impl = DefaultState(keyby)
         self._state_type = StateType.DEFAULT
 

csp_gateway/server/gateway/csp/state.py

@@ -6,13 +6,12 @@
 from datetime import datetime, timedelta
 from socket import gethostname
 from time import sleep
-from typing import Any, Callable, List, Optional, Type, Union, get_args, get_origin
+from typing import Any, Callable, List, Optional, Type, Union
 
 import csp
 from csp import ts
 from pydantic import Field, PrivateAttr, create_model, model_validator
 
-from csp_gateway.server.gateway import State
 from csp_gateway.server.settings import Settings
 
 from .csp import Channels, ChannelsFactory, ChannelsType, Module
@@ -113,19 +112,13 @@ def _instantiate_dynamic_channel(self, modules: List[Module[GatewayChannels]], c
         for m in modules:
             module_dynamic_channels = m.dynamic_channels() if hasattr(m, "dynamic_channels") else None
             if module_dynamic_channels:
-                channels_with_state = m.dynamic_state_channels() if hasattr(m, "dynamic_state_channels") else None
                 for n, t in module_dynamic_channels.items():
                     existing_type = dynamic_channels.get(n, None)
                     if existing_type is not None:
                         if t is not existing_type:
                             raise ValueError(f"Conflicting types for dynamic channel {n}.")
 
                     dynamic_channels[n] = t
-                    if channels_with_state and n in channels_with_state:
-                        if get_origin(t) is list:
-                            t = get_args(t)[0]
-
-                        dynamic_channels[f"s_{n}"] = State[t]
 
         if dynamic_channels:
             dynamic_channel_kwargs = {n: (ts[t], None) for n, t in dynamic_channels.items()}

csp_gateway/server/gateway/gateway.py

@@ -383,7 +383,7 @@ def _convert_log_level(cls, v: Union[str, int]) -> int:
     def connect(self, channels: ChannelsType):
         logger_to_use = logging.getLogger(self.log_name)
 
-        for field in self.selection.select_from(channels, state_channels=self.log_states):
+        for field in self.selection.select_from(channels):
             data = channels.get_channel(field)
             # list baskets not supported yet
             if isinstance(data, dict):
@@ -393,6 +393,11 @@ def connect(self, channels: ChannelsType):
                 edge = channels.get_channel(field)
                 csp.log(self.log_level, field, edge, logger=logger_to_use)
 
+        if self.log_states:
+            for alias in self.selection.select_from(channels, state_channels=True):
+                edge = channels.get_state(alias)
+                csp.log(self.log_level, f"s_{alias}", edge, logger=logger_to_use)
+
 
 # Backwards compatibility alias
 StdlibLogging = Logging

csp_gateway/server/modules/logging/stdlib.py

@@ -91,7 +91,8 @@ def _mount_send(self, app: GatewayWebApp) -> None:
 
     def _mount_state(self, app: GatewayWebApp) -> None:
         selection = ChannelSelection() if self.force_mount_all else self.mount_state
-        channels_set = set(selection.select_from(app.gateway.channels_model, state_channels=True))
+        # Use the instance so dynamically-registered set_state aliases are included.
+        channels_set = set(selection.select_from(app.gateway.channels, state_channels=True))
         seen_channels = set()
 
         # Bind every wire
@@ -110,7 +111,7 @@ def _mount_state(self, app: GatewayWebApp) -> None:
 
             missing_channels = channels_set - seen_channels
             if missing_channels:
-                log.info(f"Requested channels missing state routes: {list(channel[2:] for channel in missing_channels)}")
+                log.info(f"Requested channels missing state routes: {list(missing_channels)}")
 
         app.add_state_available_channels(seen_channels)