Refactor message handling to remove GeneratorMessage and streamline session management.

Author: grongierisc

demo/python/generator/msg.py

@@ -1,8 +1,8 @@
-from iop import GeneratorMessage,Message
+from iop import Message
 from dataclasses import dataclass
 
 @dataclass
-class MyGenerator(GeneratorMessage):
+class MyGenerator(Message):
     """Base message to initialize generator function"""
     my_string: str
 

src/iop/__init__.py

@@ -3,7 +3,7 @@
 from iop._business_service import _BusinessService
 from iop._director import _Director
 from iop._inbound_adapter import _InboundAdapter
-from iop._message import _Message, _PickleMessage, _PydanticMessage, _PydanticPickleMessage, _GeneratorMessage
+from iop._message import _Message, _PickleMessage, _PydanticMessage, _PydanticPickleMessage
 from iop._outbound_adapter import _OutboundAdapter
 from iop._private_session_duplex import _PrivateSessionDuplex
 from iop._private_session_process import _PrivateSessionProcess
@@ -22,5 +22,4 @@ class Message(_Message): pass
 class PickleMessage(_PickleMessage): pass
 class PydanticMessage(_PydanticMessage): pass
 class PydanticPickleMessage(_PydanticPickleMessage): pass
-class GeneratorMessage(_GeneratorMessage): pass
 class Director(_Director): pass

src/iop/_business_host.py

@@ -4,8 +4,8 @@
 from . import _iris
 from ._common import _Common
 from ._message import _Message as Message
-from ._decorators import input_serializer_param, output_deserializer
-from ._dispatch import dispatch_serializer, dispatch_deserializer
+from ._decorators import input_serializer_param, output_deserializer, input_deserializer, output_serializer
+from ._dispatch import dispatch_serializer, dispatch_deserializer, dispach_message
 from ._async_request import AsyncRequest
 from ._generator_request import _GeneratorRequest
 
@@ -68,8 +68,7 @@ async def send_request_async_ng(self, target: str, request: Union[Message, Any],
             Response from target component
         """
         return await AsyncRequest(target, request, timeout, description, self)
-    
-    @input_serializer_param(1, 'request')
+
     def send_generator_request(self, target: str, request: Union[Message, Any], 
                               timeout: int = -1, description: Optional[str] = None) -> _GeneratorRequest:
         """Send message as a generator request to target component.
@@ -81,7 +80,7 @@ def send_generator_request(self, target: str, request: Union[Message, Any],
         Returns:
             _GeneratorRequest: An instance of _GeneratorRequest to iterate over responses
         Raises:
-            TypeError: If request is not of type Message or _GeneratorMessage
+            TypeError: If request is not of type Message
         """
         return _GeneratorRequest(self, target, request, timeout, description)
 
@@ -196,7 +195,7 @@ def on_get_connections(self) -> List[str]:
         """
         ## Parse the class code to find all invocations of send_request_sync and send_request_async
         ## and return the targets
-        targer_list = []
+        target_list = []
         # get the source code of the class
         source = getsource(self.__class__)
         # find all invocations of send_request_sync and send_request_async
@@ -211,29 +210,47 @@ def on_get_connections(self) -> List[str]:
                         if target.find("=") != -1:
                             # it's a keyword argument, remove the keyword
                             target = target[target.find("=")+1:].strip()
-                        if target not in targer_list:
-                            targer_list.append(target)
+                        if target not in target_list:
+                            target_list.append(target)
                 i = source.find(method, i+1)
 
-        for target in targer_list:
+        for target in target_list:
             # if target is a string, remove the quotes
             if target[0] == "'" and target[-1] == "'":
-                targer_list[targer_list.index(target)] = target[1:-1]
+                target_list[target_list.index(target)] = target[1:-1]
             elif target[0] == '"' and target[-1] == '"':
-                targer_list[targer_list.index(target)] = target[1:-1]
+                target_list[target_list.index(target)] = target[1:-1]
             # if target is a variable, try to find the value of the variable
             else:
                 self.on_init()
                 try:
                     if target.find("self.") != -1:
                         # it's a class variable
-                        targer_list[targer_list.index(target)] = getattr(self, target[target.find(".")+1:])
+                        target_list[target_list.index(target)] = getattr(self, target[target.find(".")+1:])
                     elif target.find(".") != -1:
                         # it's a class variable
-                        targer_list[targer_list.index(target)] = getattr(getattr(self, target[:target.find(".")]), target[target.find(".")+1:])
+                        target_list[target_list.index(target)] = getattr(getattr(self, target[:target.find(".")]), target[target.find(".")+1:])
                     else:
-                        targer_list[targer_list.index(target)] = getattr(self, target)
+                        target_list[target_list.index(target)] = getattr(self, target)
                 except Exception as e:
                     pass
 
-        return targer_list
+        return target_list
+
+    @input_deserializer
+    def _dispatch_generator_started(self, request: Any) -> Any:
+        """For internal use only."""
+        self._gen = dispach_message(self, request)
+        # check if self._gen is a generator
+        if not hasattr(self._gen, '__iter__'):
+            raise TypeError("Expected a generator or iterable object, got: {}".format(type(self._gen).__name__))
+        
+        return _iris.get_iris().IOP.Generator.Message.Ack._New()
+    
+    @output_serializer
+    def _dispatch_generator_poll(self) -> Any:
+        """For internal use only."""
+        try:
+            return next(self._gen)
+        except StopIteration:
+            return _iris.get_iris().IOP.Generator.Message.Stop._New()

src/iop/_business_operation.py

@@ -73,33 +73,3 @@ def OnMessage(self, request: Any) -> Any:
         """
         return
 
-    @input_deserializer
-    @output_serializer
-    def _dispatch_private_session_started(self, request: Any) -> Any:
-        """For internal use only."""
-        self._gen = self.on_private_session_started(request)
-        from . import _iris
-        return _iris.get_iris().IOP.PrivateSession.Message.Ack._New()
-    
-    def on_private_session_started(self, request: Any) -> Any:
-        """Called when a private session is started.
-        
-        This method can be overridden to handle the start of a private session.
-        
-        Args:
-            request: The request object containing session details.
-            
-        Returns:
-            Any response object or None.
-        """
-        return None
-    
-    @input_serializer
-    @output_deserializer
-    def _dispatch_private_session_poll(self, request: Any) -> Any:
-        """For internal use only."""
-        try:
-            return next(self._gen)
-        except StopIteration:
-            import _iris
-            return _iris.get_iris().IOP.PrivateSession.Message.Stop._New()

src/iop/_dispatch.py

@@ -1,12 +1,10 @@
 from inspect import signature, Parameter
 from typing import Any, List, Tuple, Callable
 
-from ._serialization import serialize_message, serialize_pickle_message, deserialize_message, deserialize_pickle_message
+from ._serialization import serialize_message, serialize_pickle_message, deserialize_message, deserialize_pickle_message, serialize_message_generator, serialize_pickle_message_generator
 from ._message_validator import is_message_instance, is_pickle_message_instance, is_iris_object_instance
 
-from ._message import _GeneratorMessage
-
-def dispatch_serializer(message: Any) -> Any:
+def dispatch_serializer(message: Any, is_generator: bool = False) -> Any:
     """Serializes the message based on its type.
     
     Args:
@@ -20,8 +18,12 @@ def dispatch_serializer(message: Any) -> Any:
     """
     if message is not None:
         if is_message_instance(message):
+            if is_generator:
+                return serialize_message_generator(message)
             return serialize_message(message)
         elif is_pickle_message_instance(message):
+            if is_generator:
+                return serialize_pickle_message_generator(message)
             return serialize_pickle_message(message)
         elif is_iris_object_instance(message):
             return message
@@ -75,9 +77,6 @@ def dispach_message(host: Any, request: Any) -> Any:
     module = request.__class__.__module__
     classname = request.__class__.__name__
 
-    if isinstance(request, _GeneratorMessage):
-        return getattr(host, '_dispatch_private_session_started')(request)
-
     for msg, method in host.DISPATCH:
         if msg == module + "." + classname:
             call = method

src/iop/_generator_request.py

@@ -1,36 +1,30 @@
 from typing import Any, Optional, Union
 
 from . import _iris
-from ._message import _GeneratorMessage
-
+from ._dispatch import dispatch_serializer, dispatch_deserializer
 
 class _GeneratorRequest:
     """Generator class to interetate over responses from a request.
     This class is used to handle the responses from a request in a generator-like manner."""
 
-    def __init__(self, host: Any, target: str, request: Union[_GeneratorMessage, Any], 
+    def __init__(self, host: Any, target: str, request: Any, 
                  timeout: int = -1, description: Optional[str] = None) -> None:
         self.host = host
         self.target = target
         self.request = request
-        self.timeout = timeout
-        self.description = description
-        self._response = None
 
-        # if not isinstance(self.request, _GeneratorMessage):
-        #     raise TypeError("request must be of type Message or _GeneratorMessage")
-        
-        ack_rsponse = self.host.send_request_sync(self.target, self.request)
+        ack_response = self.host.send_request_sync(self.target, dispatch_serializer(self.request, is_generator=True),
+                                                  timeout=timeout, description=description)
 
-        if ack_rsponse is None:
+        if ack_response is None or not ack_response._IsA("IOP.Generator.Message.Ack"):
             raise RuntimeError("Failed to send request, no acknowledgment received.")
 
     def __iter__(self):
         return self
 
     def __next__(self):
-        poll = _iris.get_iris().IOP.PrivateSession.Message.Poll._New()
+        poll = _iris.get_iris().IOP.Generator.Message.Poll._New()
         rsp = self.host.send_request_sync(self.target, poll)
-        if rsp is None:
+        if rsp is None or (hasattr(rsp, '_IsA') and rsp._IsA("IOP.Generator.Message.Stop")):
             raise StopIteration("No more responses available.")
         return rsp

src/iop/_message.py

@@ -9,7 +9,7 @@ class _Message:
     """
     pass
 
-class _PickleMessage:
+class _PickleMessage(_Message):
     """ The abstract class that is the superclass for persistent messages sent from one component to another.
     This class has no properties or methods. Users subclass Message and add properties.
     The IOP framework provides the persistence to objects derived from the Message class.
@@ -28,5 +28,3 @@ class _PydanticPickleMessage(BaseModel):
     def __init__(self, **data: Any):
         super().__init__(**data)
 
-class _GeneratorMessage(_Message):
-    """"Base message to initialize generator function"""

src/iop/_serialization.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel, TypeAdapter, ValidationError
 
 from . import _iris
-from ._message import _PydanticPickleMessage, _Message, _GeneratorMessage
+from ._message import _PydanticPickleMessage, _Message
 from ._utils import _Utils
 
 class SerializationError(Exception):
@@ -37,18 +37,18 @@ def _convert_to_json_safe(obj: Any) -> Any:
             raise SerializationError(f"Object {obj} must be a Pydantic model or dataclass Message")
 
     @staticmethod
-    def serialize(message: Any, use_pickle: bool = False) -> Any:
+    def serialize(message: Any, use_pickle: bool = False, is_generator:bool = False) -> Any:
         """Serializes a message to IRIS format."""
-        if isinstance(message, _PydanticPickleMessage) or use_pickle:
-            return MessageSerializer._serialize_pickle(message)
-        return MessageSerializer._serialize_json(message)
+        if use_pickle:
+            return MessageSerializer._serialize_pickle(message, is_generator)
+        return MessageSerializer._serialize_json(message, is_generator)
 
     @staticmethod
-    def _serialize_json(message: Any) -> Any:
+    def _serialize_json(message: Any, is_generator: bool = False) -> Any:
         json_string = MessageSerializer._convert_to_json_safe(message)
 
-        if isinstance(json_string, _GeneratorMessage):
-            msg = _iris.get_iris().cls('IOP.PrivateSession.Message.Start')._New()
+        if is_generator:
+            msg = _iris.get_iris().cls('IOP.Generator.Message.Start')._New()
         else:
             msg = _iris.get_iris().cls('IOP.Message')._New()
         msg.classname = f"{message.__class__.__module__}.{message.__class__.__name__}"
@@ -91,9 +91,12 @@ def _deserialize_json(serial: Any) -> Any:
             raise SerializationError(f"Failed to deserialize JSON: {str(e)}")
 
     @staticmethod
-    def _serialize_pickle(message: Any) -> Any:
+    def _serialize_pickle(message: Any, is_generator: bool = False) -> Any:
         pickle_string = codecs.encode(pickle.dumps(message), "base64").decode()
-        msg = _iris.get_iris().cls('IOP.PickleMessage')._New()
+        if is_generator:
+            msg = _iris.get_iris().cls('IOP.Generator.Message.StartPickle')._New()
+        else:
+            msg = _iris.get_iris().cls('IOP.PickleMessage')._New()
         msg.classname = f"{message.__class__.__module__}.{message.__class__.__name__}"
         msg.jstr = _Utils.string_to_stream(pickle_string)
         return msg
@@ -168,7 +171,9 @@ def dataclass_to_dict(instance: Any) -> Dict:
     return result
 
 # Maintain backwards compatibility
-serialize_pickle_message = lambda msg: MessageSerializer.serialize(msg, use_pickle=True)
-serialize_message = lambda msg: MessageSerializer.serialize(msg, use_pickle=False)
+serialize_pickle_message = lambda msg: MessageSerializer.serialize(msg, use_pickle=True, is_generator=False)
+serialize_pickle_message_generator = lambda msg: MessageSerializer.serialize(msg, use_pickle=True, is_generator=True)
+serialize_message = lambda msg: MessageSerializer.serialize(msg, use_pickle=False, is_generator=False)
+serialize_message_generator = lambda msg: MessageSerializer.serialize(msg, use_pickle=False, is_generator=True)
 deserialize_pickle_message = lambda serial: MessageSerializer.deserialize(serial, use_pickle=True)
 deserialize_message = lambda serial: MessageSerializer.deserialize(serial, use_pickle=False)