fix: update imports

This commit is contained in:
2023-12-07 16:17:59 +01:00
parent 95efd78c64
commit e3193f2cc3
8 changed files with 26 additions and 38 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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}/<path_to_my_config_file.yaml>"

View File

@@ -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.

View File

@@ -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:

View File

@@ -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:

View File

@@ -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)

View File

@@ -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))}
),
]