ezmsg-baseproc 1.2.1__tar.gz → 1.3.0__tar.gz

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-baseproc
-Version: 1.2.1
+Version: 1.3.0
 Summary: Base processor classes and protocols for ezmsg signal processing pipelines
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/src/ezmsg/baseproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.2.1'
-__version_tuple__ = version_tuple = (1, 2, 1)
+__version__ = version = '1.3.0'
+__version_tuple__ = version_tuple = (1, 3, 0)
 
 __commit_id__ = commit_id = None

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/src/ezmsg/baseproc/protocols.py

@@ -5,8 +5,7 @@ import typing
 from dataclasses import dataclass
 
 import ezmsg.core as ez
-
-from .util.message import SampleMessage
+from ezmsg.util.messages.axisarray import AxisArray
 
 # --- Processor state decorator ---
 processor_state = functools.partial(dataclass, unsafe_hash=True, frozen=False, init=False)
@@ -138,7 +137,7 @@ class StatefulTransformer(
 
 
 class AdaptiveTransformer(StatefulTransformer, typing.Protocol):
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         """Update transformer state using labeled training data.
 
         This method should update the internal state/parameters of the transformer
@@ -146,4 +145,4 @@ class AdaptiveTransformer(StatefulTransformer, typing.Protocol):
         """
         ...
 
-    async def apartial_fit(self, message: SampleMessage) -> None: ...
+    async def apartial_fit(self, message: AxisArray) -> None: ...

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/src/ezmsg/baseproc/stateful.py

@@ -4,6 +4,9 @@ import pickle
 import typing
 from abc import ABC, abstractmethod
 
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
 from .processor import (
     BaseProcessor,
     BaseProducer,
@@ -256,7 +259,7 @@ class BaseStatefulTransformer(
 class BaseAdaptiveTransformer(
     BaseStatefulTransformer[
         SettingsType,
-        MessageInType | SampleMessage,
+        MessageInType,
         MessageOutType | None,
         StateType,
     ],
@@ -264,30 +267,41 @@ class BaseAdaptiveTransformer(
     typing.Generic[SettingsType, MessageInType, MessageOutType, StateType],
 ):
     @abstractmethod
-    def partial_fit(self, message: SampleMessage) -> None: ...
+    def partial_fit(self, message: AxisArray) -> None: ...
 
-    async def apartial_fit(self, message: SampleMessage) -> None:
+    async def apartial_fit(self, message: AxisArray) -> None:
         """Override me if you need async partial fitting."""
         return self.partial_fit(message)
 
-    def __call__(self, message: MessageInType | SampleMessage) -> MessageOutType | None:
+    def __call__(self, message: MessageInType) -> MessageOutType | None:
         """
         Adapt transformer with training data (and optionally labels)
-        in SampleMessage
+        in AxisArray with attrs["trigger"].
 
         Args:
-            message: An instance of SampleMessage with optional
-             labels (y) in message.trigger.value.data and
-             data (X) in message.sample.data
+            message: An AxisArray with optional trigger in attrs["trigger"],
+             containing labels (y) in attrs["trigger"].value and
+             data (X) in message.data
 
         Returns: None
         """
         if is_sample_message(message):
+            if isinstance(message, SampleMessage):
+                # Auto-convert old format → new format
+                message = replace(
+                    message.sample,
+                    attrs={**message.sample.attrs, "trigger": message.trigger},
+                )
             return self.partial_fit(message)
         return super().__call__(message)
 
-    async def __acall__(self, message: MessageInType | SampleMessage) -> MessageOutType | None:
+    async def __acall__(self, message: MessageInType) -> MessageOutType | None:
         if is_sample_message(message):
+            if isinstance(message, SampleMessage):
+                message = replace(
+                    message.sample,
+                    attrs={**message.sample.attrs, "trigger": message.trigger},
+                )
             return await self.apartial_fit(message)
         return await super().__acall__(message)
 

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/src/ezmsg/baseproc/units.py

@@ -3,10 +3,10 @@
 import math
 import traceback
 import typing
+import warnings
 from abc import ABC, abstractmethod
 
 import ezmsg.core as ez
-from ezmsg.util.generator import GenState
 from ezmsg.util.messages.axisarray import AxisArray, LinearAxis
 
 from .clockdriven import BaseClockDrivenProducer
@@ -14,7 +14,6 @@ from .composite import CompositeProcessor
 from .processor import BaseConsumer, BaseProducer, BaseTransformer
 from .protocols import MessageInType, MessageOutType, SettingsType
 from .stateful import BaseAdaptiveTransformer, BaseStatefulConsumer, BaseStatefulTransformer
-from .util.message import SampleMessage
 from .util.profile import profile_subpub
 from .util.typeresolution import resolve_typevar
 
@@ -223,7 +222,7 @@ class BaseAdaptiveTransformerUnit(
     ABC,
     typing.Generic[SettingsType, MessageInType, MessageOutType, AdaptiveTransformerType],
 ):
-    INPUT_SAMPLE = ez.InputStream(SampleMessage)
+    INPUT_SAMPLE = ez.InputStream(AxisArray)
     INPUT_SIGNAL = ez.InputStream(MessageInType)
     OUTPUT_SIGNAL = ez.OutputStream(MessageOutType)
 
@@ -242,7 +241,7 @@ class BaseAdaptiveTransformerUnit(
             yield self.OUTPUT_SIGNAL, result
 
     @ez.subscriber(INPUT_SAMPLE)
-    async def on_sample(self, msg: SampleMessage) -> None:
+    async def on_sample(self, msg: AxisArray) -> None:
         await self.processor.apartial_fit(msg)
 
 
@@ -287,10 +286,42 @@ class BaseClockDrivenUnit(
             yield self.OUTPUT_SIGNAL, result
 
 
-# Legacy class
+class GenState(ez.State):
+    """
+    .. deprecated::
+        ``GenState`` is deprecated. Define a local state class or use
+        ``ezmsg.baseproc`` processor classes instead.
+    """
+
+    gen: typing.Generator[typing.Any, typing.Any, None]
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        warnings.warn(
+            "GenState is deprecated. Define a local state class instead of subclassing GenState.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+
 class GenAxisArray(ez.Unit):
+    """
+    .. deprecated::
+        ``GenAxisArray`` is deprecated. Use ``BaseTransformerUnit`` or
+        ``BaseAdaptiveTransformerUnit`` from ``ezmsg.baseproc`` instead.
+    """
+
     STATE = GenState
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        warnings.warn(
+            "GenAxisArray is deprecated. Use BaseTransformerUnit or "
+            "BaseAdaptiveTransformerUnit from ezmsg.baseproc instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
     INPUT_SETTINGS = ez.InputStream(ez.Settings)

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/src/ezmsg/baseproc/util/message.py

@@ -1,5 +1,6 @@
 import time
 import typing
+import warnings
 from dataclasses import dataclass, field
 
 from ezmsg.util.messages.axisarray import AxisArray
@@ -19,13 +20,28 @@ class SampleTriggerMessage:
 
 @dataclass
 class SampleMessage:
+    """
+    .. deprecated::
+        ``SampleMessage`` is deprecated. Use ``AxisArray`` with
+        ``attrs={"trigger": SampleTriggerMessage(...)}`` instead.
+    """
+
     trigger: SampleTriggerMessage
     """The time, window, and value (if any) associated with the trigger."""
 
     sample: AxisArray
     """The data sampled around the trigger."""
 
+    def __post_init__(self):
+        warnings.warn(
+            "SampleMessage is deprecated. Use AxisArray with " "attrs={'trigger': SampleTriggerMessage(...)} instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
 
-def is_sample_message(message: typing.Any) -> typing.TypeGuard[SampleMessage]:
-    """Check if the message is a SampleMessage."""
-    return hasattr(message, "trigger")
+def is_sample_message(message: typing.Any) -> bool:
+    """Detect old SampleMessage OR new AxisArray-with-trigger."""
+    if isinstance(message, SampleMessage):
+        return True
+    return isinstance(message, AxisArray) and "trigger" in getattr(message, "attrs", {})

{ezmsg_baseproc-1.2.1 → ezmsg_baseproc-1.3.0}/tests/test_baseproc.py

@@ -3,11 +3,11 @@
 import dataclasses
 import pickle
 from types import NoneType
-from typing import Any, Generator
-from unittest.mock import MagicMock
+from typing import Any
 
+import numpy as np
 import pytest
-from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import AxisArray
 
 from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
@@ -22,7 +22,7 @@ from ezmsg.baseproc import (
     BaseTransformer,
     CompositeProcessor,
     CompositeProducer,
-    SampleMessage,
+    SampleTriggerMessage,
     _get_base_processor_message_in_type,
     _get_base_processor_message_out_type,
     _get_base_processor_settings_type,
@@ -136,11 +136,11 @@ class MockAdaptiveTransformer(BaseAdaptiveTransformer[MockSettings, MockMessageA
     def _reset_state(self, message: MockMessageA) -> None:
         self._state.iterations = 0
 
-    def _process(self, message: MockMessageA | SampleMessage) -> MockMessageB:
+    def _process(self, message: MockMessageA) -> MockMessageB:
         self._state.iterations += 1
         return MockMessageB()
 
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         self._state.iterations += 1
 
 
@@ -251,18 +251,18 @@ class ChainedCompositeProcessorWithDeepProcessors(CompositeProcessor[MockSetting
         }
 
 
-@consumer
-def mock_generator() -> Generator[MockMessageA, MockMessageA, None]:
-    """A mock generator function for testing purposes."""
-    while True:
-        yield MockMessageA()
+class MockGeneratorTransformer(BaseTransformer[MockSettings, MockMessageA, MockMessageA]):
+    """A mock transformer replacing the legacy @consumer generator."""
+
+    def _process(self, message: MockMessageA) -> MockMessageA:
+        return MockMessageA()
 
 
 class MockGeneratorCompositeProcessor(CompositeProcessor[MockSettings, MockMessageA, MockMessageB]):
     @staticmethod
     def _initialize_processors(settings):
         return {
-            "generator": mock_generator(),
+            "generator": MockGeneratorTransformer(settings=settings),
             "stateful_processor": MockStatefulProcessor(settings=settings),
         }
 
@@ -333,27 +333,26 @@ class ChainedCompositeProducerWithDeepProcessors(CompositeProducer[MockSettings,
         }
 
 
-@consumer
-def mock_producer_generator() -> Generator[MockMessageA, None, None]:
-    """A mock generator function for testing purposes."""
-    while True:
-        yield MockMessageA()
+class MockGeneratorProducer(BaseProducer[MockSettings, MockMessageA]):
+    """A mock producer replacing the legacy @consumer producer generator."""
+
+    async def _produce(self) -> MockMessageA:
+        return MockMessageA()
+
 
+class MockGeneratorPassthroughTransformer(BaseTransformer[MockSettings, MockMessageA, MockMessageA]):
+    """A mock transformer replacing the legacy unprimed generator."""
 
-def mock_generator_unprimed() -> Generator[MockMessageA, MockMessageA, None]:
-    """A mock generator function for testing purposes."""
-    output = MockMessageA()
-    while True:
-        input = yield output
-        output = input or output
+    def _process(self, message: MockMessageA) -> MockMessageA:
+        return message or MockMessageA()
 
 
 class MockGeneratorCompositeProducer(CompositeProducer[MockSettings, MockMessageB]):
     @staticmethod
     def _initialize_processors(settings):
         return {
-            "generator": mock_producer_generator(),
-            "mock_generator_unprimed": mock_generator_unprimed(),
+            "generator": MockGeneratorProducer(settings=settings),
+            "mock_generator_passthrough": MockGeneratorPassthroughTransformer(settings=settings),
             "stateful_processor": MockStatefulProcessor(settings=settings),
         }
 
@@ -758,10 +757,13 @@ class TestBaseStatefulTransformer:
         assert new_state[0].iterations == 1
 
 
-# Mock SampleMessage for testing BaseAdaptiveTransformer
+# Helper to create an AxisArray with trigger in attrs for testing BaseAdaptiveTransformer
 def mock_sample_message():
-    sample_message = MagicMock(spec=SampleMessage)
-    return sample_message
+    return AxisArray(
+        data=np.zeros((1, 1)),
+        dims=["time", "ch"],
+        attrs={"trigger": SampleTriggerMessage()},
+    )
 
 
 class TestBaseAdaptiveTransformer:
@@ -778,9 +780,7 @@ class TestBaseAdaptiveTransformer:
 
     def test_call_with_sample_message(self):
         transformer = MockAdaptiveTransformer()
-        # Create a sample message with a trigger attribute
         sample_msg = mock_sample_message()
-        setattr(sample_msg, "trigger", None)
         result = transformer(sample_msg)
         assert result is None  # partial_fit returns None
         assert transformer.state.iterations == 1