From e3193f2cc3ec575bd9afd077df3aec62e77d26ea Mon Sep 17 00:00:00 2001 From: Ivan Usov Date: Thu, 7 Dec 2023 16:17:59 +0100 Subject: [PATCH] fix: update imports --- bec_plugins/bec_client/hli/spec_hli.py | 4 ++-- .../bec_client/startup/post_startup.py | 2 +- bec_plugins/bec_client/startup/pre_startup.py | 2 +- .../services/NIDAQ_writer/NIDAQ_writer.py | 24 +++++++------------ .../services/NIDAQ_writer/scan_status.py | 6 ++--- .../services/NIDAQ_writer/utils/nidaq_sim.py | 12 +++------- bec_plugins/services/launch_writer.py | 2 +- tests/test_file_writer_service.py | 12 +++++----- 8 files changed, 26 insertions(+), 38 deletions(-) diff --git a/bec_plugins/bec_client/hli/spec_hli.py b/bec_plugins/bec_client/hli/spec_hli.py index 5d699c2..79700b2 100644 --- a/bec_plugins/bec_client/hli/spec_hli.py +++ b/bec_plugins/bec_client/hli/spec_hli.py @@ -1,5 +1,5 @@ -from bec_client.scan_manager import ScanReport -from bec_utils.devicemanager import Device +from bec_lib.devicemanager import Device +from bec_lib.scan_report import ScanReport # pylint:disable=undefined-variable # pylint: disable=too-many-arguments diff --git a/bec_plugins/bec_client/startup/post_startup.py b/bec_plugins/bec_client/startup/post_startup.py index a0f805f..3314862 100644 --- a/bec_plugins/bec_client/startup/post_startup.py +++ b/bec_plugins/bec_client/startup/post_startup.py @@ -25,7 +25,7 @@ parse the --session argument, add the following lines to the script: # pylint: disable=invalid-name, unused-import, import-error, undefined-variable, unused-variable, unused-argument, no-name-in-module import argparse -from bec_lib.core import bec_logger +from bec_lib import bec_logger logger = bec_logger.logger diff --git a/bec_plugins/bec_client/startup/pre_startup.py b/bec_plugins/bec_client/startup/pre_startup.py index dcfa194..b39c6eb 100644 --- a/bec_plugins/bec_client/startup/pre_startup.py +++ b/bec_plugins/bec_client/startup/pre_startup.py @@ -8,7 +8,7 @@ To set up the BEC client configuration, use the ServiceConfig class. For example to set the configuration file path, add the following lines to the script: import pathlib - from bec_lib.core import ServiceConfig + from bec_lib import ServiceConfig current_path = pathlib.Path(__file__).parent.resolve() CONFIG_PATH = f"{current_path}/" diff --git a/bec_plugins/services/NIDAQ_writer/NIDAQ_writer.py b/bec_plugins/services/NIDAQ_writer/NIDAQ_writer.py index a0f5694..dfdd40a 100644 --- a/bec_plugins/services/NIDAQ_writer/NIDAQ_writer.py +++ b/bec_plugins/services/NIDAQ_writer/NIDAQ_writer.py @@ -8,16 +8,10 @@ from typing import List import h5py import numpy as np -from bec_lib.core import ( - BECMessage, - MessageEndpoints, - RedisConnector, - ServiceConfig, - bec_logger, -) -from bec_lib.core.bec_service import BECService -from bec_lib.core.file_utils import FileWriterMixin -from bec_lib.core.redis_connector import MessageObject +from bec_lib import MessageEndpoints, RedisConnector, ServiceConfig, bec_logger, messages +from bec_lib.bec_service import BECService +from bec_lib.file_utils import FileWriterMixin +from bec_lib.redis_connector import MessageObject logger = bec_logger.logger @@ -65,7 +59,7 @@ class NIDAQWriterService(BECService): """ Callback for scan status messages. """ - msg = BECMessage.ScanStatusMessage.loads(message.value) + msg = messages.ScanStatusMessage.loads(message.value) if not msg: return parent.handle_scan_status(msg) @@ -114,7 +108,7 @@ class NIDAQWriterService(BECService): logger.debug( f"Received {num_msgs} messages in {time.time() - start_time} seconds" ) - msgs = [BECMessage.DeviceMessage.loads(m[1][b"device_msg"]) for m in msg[0][1]] + msgs = [messages.DeviceMessage.loads(m[1][b"device_msg"]) for m in msg[0][1]] start_time = time.time() self.handle_ni_data(msgs) logger.debug( @@ -129,7 +123,7 @@ class NIDAQWriterService(BECService): time.sleep(0.001) if msgs: try: - msgs = [BECMessage.DeviceMessage.loads(msg) for msg in msgs] + msgs = [messages.DeviceMessage.loads(msg) for msg in msgs] print(f"Received {len(msgs)} messages in {time.time() - start_time} seconds") start_time = time.time() self.handle_ni_data(msgs) @@ -138,7 +132,7 @@ class NIDAQWriterService(BECService): content = traceback.format_exc() logger.error(f"Failed to parse message: {content}") - def handle_scan_status(self, msg: BECMessage.ScanStatusMessage) -> None: + def handle_scan_status(self, msg: messages.ScanStatusMessage) -> None: """ Handle scan status messages. @@ -153,7 +147,7 @@ class NIDAQWriterService(BECService): else: self.scan_is_running = False - def handle_ni_data(self, msgs: List[BECMessage.DeviceMessage]) -> None: + def handle_ni_data(self, msgs: List[messages.DeviceMessage]) -> None: """ Receive NI data messages and write them to the writer queue. diff --git a/bec_plugins/services/NIDAQ_writer/scan_status.py b/bec_plugins/services/NIDAQ_writer/scan_status.py index f50748b..ac53d87 100644 --- a/bec_plugins/services/NIDAQ_writer/scan_status.py +++ b/bec_plugins/services/NIDAQ_writer/scan_status.py @@ -1,13 +1,13 @@ -from bec_lib.core import BECMessage, MessageEndpoints, RedisConnector +from bec_lib import MessageEndpoints, RedisConnector, messages def send_scan_status(scan_number, status): if status == "start": - scan_status_msg = BECMessage.ScanStatusMessage( + scan_status_msg = messages.ScanStatusMessage( scanID="test", status="open", info={"scan_number": scan_number} ) elif status == "stop": - scan_status_msg = BECMessage.ScanStatusMessage( + scan_status_msg = messages.ScanStatusMessage( scanID="test", status="closed", info={"scan_number": scan_number} ) else: diff --git a/bec_plugins/services/NIDAQ_writer/utils/nidaq_sim.py b/bec_plugins/services/NIDAQ_writer/utils/nidaq_sim.py index 56835a7..6b4fe15 100644 --- a/bec_plugins/services/NIDAQ_writer/utils/nidaq_sim.py +++ b/bec_plugins/services/NIDAQ_writer/utils/nidaq_sim.py @@ -2,13 +2,7 @@ import threading import time import numpy as np -from bec_lib.core import ( - BECMessage, - MessageEndpoints, - RedisConnector, - ServiceConfig, - bec_logger, -) +from bec_lib import MessageEndpoints, RedisConnector, ServiceConfig, bec_logger, messages class NIDAQSim(threading.Thread): @@ -24,9 +18,9 @@ class NIDAQSim(threading.Thread): "signal2": signal, } - msg = BECMessage.DeviceMessage(signals=signals) + msg = messages.DeviceMessage(signals=signals) msg = msg.dumps() - BECMessage.DeviceMessage.loads(msg) + messages.DeviceMessage.loads(msg) total_time = time.time() while True: # if index > 1000: diff --git a/bec_plugins/services/launch_writer.py b/bec_plugins/services/launch_writer.py index 4c8dc6a..f2e1099 100644 --- a/bec_plugins/services/launch_writer.py +++ b/bec_plugins/services/launch_writer.py @@ -1,7 +1,7 @@ import argparse import threading -from bec_lib.core import RedisConnector, ServiceConfig, bec_logger +from bec_lib import RedisConnector, ServiceConfig, bec_logger from NIDAQ_writer import NIDAQWriterService parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) diff --git a/tests/test_file_writer_service.py b/tests/test_file_writer_service.py index 1b1ab33..a5e838d 100644 --- a/tests/test_file_writer_service.py +++ b/tests/test_file_writer_service.py @@ -2,8 +2,8 @@ from unittest import mock import numpy as np import pytest -from bec_lib.core import BECMessage, MessageEndpoints, ServiceConfig -from bec_lib.core.redis_connector import MessageObject +from bec_lib import MessageEndpoints, ServiceConfig, messages +from bec_lib.redis_connector import MessageObject from bec_plugins.services.NIDAQ_writer import NIDAQWriterService @@ -46,7 +46,7 @@ def test_nidaq_scan_status_consumer(nidaq): def test_scan_status_callback(nidaq): - scan_status_msg = BECMessage.ScanStatusMessage(scanID="test", status="open", info={}) + scan_status_msg = messages.ScanStatusMessage(scanID="test", status="open", info={}) msg_obj = MessageObject( topic="test", value=scan_status_msg.dumps(), @@ -85,7 +85,7 @@ def test_nidaq_reads_data_from_strea(nidaq): @pytest.mark.parametrize("scan_status", ["open", "closed", "aborted", "halted", None]) def test_nidaq_handle_scan_status(nidaq, scan_status): - scan_status_msg = BECMessage.ScanStatusMessage( + scan_status_msg = messages.ScanStatusMessage( scanID="test", status=scan_status, info={"scan_number": 5} ) nidaq.handle_scan_status(scan_status_msg) @@ -99,8 +99,8 @@ def test_nidaq_handle_scan_status(nidaq, scan_status): def test_nidaq_handle_ni_data(nidaq): data = [ - BECMessage.DeviceMessage(signals={"signal1": list(range(10)), "signal2": list(range(10))}), - BECMessage.DeviceMessage( + messages.DeviceMessage(signals={"signal1": list(range(10)), "signal2": list(range(10))}), + messages.DeviceMessage( signals={"signal1": list(range(10, 20)), "signal2": list(range(10, 20))} ), ]