wrap observations in a MessageWrapper

Author: tipatterson-dev

oshconnect/datasource/datasource.py

@@ -3,10 +3,13 @@
 #  Author: Ian Patterson <ian@botts-inc.com>
 #
 #  Contact Email: ian@botts-inc.com
+from __future__ import annotations
 import asyncio
+import json
 from uuid import uuid4
 
 import websockets
+from conSys4Py.datamodels.observations import ObservationOMJSONInline
 
 from external_models.object_models import DatastreamResource
 from oshconnect import Utilities
@@ -60,37 +63,55 @@ def update_properties(self, properties: dict):
         # TODO: need to stop in progress sub-processes and restart
         self.properties = properties
 
+    def set_mode(self, mode: str):
+        self.mode = mode
+
     def initialize(self):
-        pass
+        if self._websocket.is_open():
+            self._websocket.close()
+        self._websocket = None
+        self._status = "initialized"
 
     async def connect(self):
         if self.mode == "websocket":
             self._websocket = await websockets.connect(self._url, extra_headers=self._extra_headers)
             self._status = "connected"
             return self._websocket
+        elif self.mode == "playback":
+            self._status = "connected"
+            return "Playback mode is not yet implemented."
+        elif self.mode == "live-batch":
+            self._status = "connected"
+            return "Live-batch mode is not yet implemented."
 
+    def disconnect(self):
+        self._websocket.close()
 
-def disconnect(self):
-    pass
-
-
-def reset(self):
-    pass
+    def reset(self):
+        self._websocket = None
+        self._status = "initialized"
 
+    def get_status(self):
+        return self._status
 
-def get_status(self):
-    return self.status
+    def get_ws_client(self):
+        return self._websocket
 
 
 class DataSourceHandler:
     datasource_map: dict[str, DataSource]
+    _message_list: MessageHandler
 
     def __init__(self):
         self.datasource_map = {}
+        self._message_list = MessageHandler()
 
     def add_datasource(self, datasource: DataSource):
         self.datasource_map[datasource.get_id()] = datasource
 
+    def remove_datasource(self, datasource_id: str):
+        return self.datasource_map.pop(datasource_id)
+
     def initialize_ds(self, datasource_id: str, properties: dict):
         ds = self.datasource_map.get(datasource_id)
         ds.initialize()
@@ -99,17 +120,76 @@ def initialize_all(self):
         # list comp is faster than for loop
         [ds.initialize() for ds in self.datasource_map.values()]
 
+    def set_ds_mode(self, mode: str):
+        var = (ds.set_mode(mode) for ds in self.datasource_map.values())
+
     async def connect_ds(self, datasource_id: str):
         ds = self.datasource_map.get(datasource_id)
         await ds.connect()
 
     async def connect_all(self):
-        results = await asyncio.gather(*(ds.connect() for ds in self.datasource_map.values()))
-        return results
+        # call connect for all datasources
+        [(ds, await ds.connect()) for ds in self.datasource_map.values()]
+        for ds in self.datasource_map.values():
+            task = asyncio.create_task(self._handle_datastream_client(ds))
+            # return task
 
     def disconnect_ds(self, datasource_id: str):
         ds = self.datasource_map.get(datasource_id)
         ds.disconnect()
 
     def disconnect_all(self):
         [ds.disconnect() for ds in self.datasource_map.values()]
+
+    async def _handle_datastream_client(self, datasource: DataSource):
+        try:
+            async for msg in datasource.get_ws_client():
+                msg_dict = json.loads(msg.decode('utf-8'))
+                obs = ObservationOMJSONInline.model_validate(msg_dict)
+                msg_wrapper = MessageWrapper(datasource=datasource, message=obs)
+                self._message_list.add_message(msg_wrapper)
+
+        except Exception as e:
+            print(f"An error occurred while reading from websocket: {e}")
+
+
+class MessageHandler:
+    _message_list: list[MessageWrapper]
+
+    def __init__(self):
+        self._message_list = []
+
+    def add_message(self, message: MessageWrapper):
+        self._message_list.append(message)
+        print(self._message_list)
+
+    def get_messages(self):
+        return self._message_list
+
+    def clear_messages(self):
+        self._message_list.clear()
+
+    def sort_messages(self):
+        # copy the list
+        sorted_list = self._message_list.copy()
+        sorted_list.sort(key=lambda x: x.resultTime)
+        return sorted_list
+
+
+class MessageWrapper:
+    """
+    Combines a DataSource and a Message into a single object for easier access
+    """
+
+    def __init__(self, datasource: DataSource, message: ObservationOMJSONInline):
+        self._message = message
+        self._datasource = datasource
+
+    def get_message(self):
+        return self._message
+
+    def get_message_as_dict(self):
+        return self._message.dict()
+
+    def __repr__(self):
+        return f"{self._datasource}, {self._message}"

oshconnect/oshconnect.py

@@ -82,16 +82,10 @@ def select_temporal_mode(self, mode: str):
 
     async def playback_streams(self, stream_ids: list = None):
         if stream_ids is None:
-            clients = await self._datasource_handler.connect_all()
-            for client in clients:
-                task = asyncio.create_task(self._handle_datastream_client(client))
-                self._tasks.append(task)
+            await self._datasource_handler.connect_all()
         else:
             for stream_id in stream_ids:
-                clients = await self._datasource_handler.connect_ds(stream_id)
-                for client in clients:
-                    msg = await client.recv()
-                    print(msg)
+                await self._datasource_handler.connect_ds(stream_id)
 
     def visualize_streams(self, streams: list):
         pass
@@ -128,12 +122,3 @@ def authenticate_user(self, user: dict):
 
     def synchronize_streams(self, systems: list):
         pass
-
-    async def _handle_datastream_client(self, client):
-        try:
-            async for msg in client:
-                msg_dict = json.loads(msg.decode('utf-8'))
-                obs = ObservationOMJSONInline.model_validate(msg_dict)
-
-        except Exception as e:
-            print(f"An error occurred while reading from websocket: {e}")