diff --git a/src/genie_python/genie_advanced.py b/src/genie_python/genie_advanced.py index b90cbb6..856f54f 100644 --- a/src/genie_python/genie_advanced.py +++ b/src/genie_python/genie_advanced.py @@ -5,10 +5,11 @@ """ import contextlib -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from time import sleep -from typing import Any, Callable, Iterator, TypedDict +from typing import Any, Iterator, Protocol, TypedDict +import numpy as np import numpy.typing as npt from genie_python.genie_api_setup import ( @@ -21,6 +22,11 @@ from genie_python.utilities import check_break +class PrePostCmd(Protocol): + def __call__(self, **kwargs: Any) -> str | None: + pass + + @usercommand @helparglist("") def get_manager_mode() -> bool: @@ -52,7 +58,7 @@ def assert_in_manager_mode() -> None: @contextlib.contextmanager -def motor_in_set_mode(pv_name: str) -> Iterator: +def motor_in_set_mode(pv_name: str) -> Iterator[None]: """ Uses a context to place motor into set mode and ensure that it leaves set mode after context has ended. If it can not set the mode correctly @@ -134,7 +140,7 @@ def get_pv_from_block(block: str) -> str: @usercommand @helparglist("pv str") @log_command_and_handle_exception -def pv_exists(pv: str, is_local: bool = False) -> None: +def pv_exists(pv: str, is_local: bool = False) -> bool: """ Check if PV exists. @@ -159,13 +165,13 @@ def wait_for_pv( value: The value to wait for maxwait (int, optional): The maximum time to wait for in seconds """ - start_time = datetime.utcnow() + start_time = datetime.now(UTC) while True: curr_value = __api.get_pv_value(pv) if curr_value == value: break if maxwait is not None: - if timedelta(seconds=maxwait) < datetime.utcnow() - start_time: + if timedelta(seconds=maxwait) < datetime.now(UTC) - start_time: break sleep(DELAY_IN_WAIT_FOR_SLEEP_LOOP) check_break(2) @@ -173,7 +179,7 @@ def wait_for_pv( @usercommand @helparglist("") -def set_begin_precmd(begin_precmd: Callable[[Any], str | None]) -> None: +def set_begin_precmd(begin_precmd: PrePostCmd) -> None: """ Set the function to call before the begin command. @@ -186,7 +192,7 @@ def set_begin_precmd(begin_precmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_begin_postcmd(begin_postcmd: Callable[[Any], str | None]) -> None: +def set_begin_postcmd(begin_postcmd: PrePostCmd) -> None: """ Set the function to call after the begin command. @@ -198,7 +204,7 @@ def set_begin_postcmd(begin_postcmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_abort_precmd(abort_precmd: Callable[[Any], str | None]) -> None: +def set_abort_precmd(abort_precmd: PrePostCmd) -> None: """ Set the function to call before the abort command. @@ -210,7 +216,7 @@ def set_abort_precmd(abort_precmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_abort_postcmd(abort_postcmd: Callable[[Any], str | None]) -> None: +def set_abort_postcmd(abort_postcmd: PrePostCmd) -> None: """ Set the function to call after the abort command. @@ -222,7 +228,7 @@ def set_abort_postcmd(abort_postcmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_end_precmd(end_precmd: Callable[[Any], str | None]) -> None: +def set_end_precmd(end_precmd: PrePostCmd) -> None: """ Set the function to call before the end command. @@ -234,7 +240,7 @@ def set_end_precmd(end_precmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_end_postcmd(end_postcmd: Callable[[Any], str | None]) -> None: +def set_end_postcmd(end_postcmd: PrePostCmd) -> None: """ Set the function to call after the end command. @@ -246,7 +252,7 @@ def set_end_postcmd(end_postcmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_pause_precmd(pause_precmd: Callable[[Any], str | None]) -> None: +def set_pause_precmd(pause_precmd: PrePostCmd) -> None: """ Set the function to call before the pause command. @@ -258,7 +264,7 @@ def set_pause_precmd(pause_precmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_pause_postcmd(pause_postcmd: Callable[[Any], str | None]) -> None: +def set_pause_postcmd(pause_postcmd: PrePostCmd) -> None: """ Set the function to call after the pause command. @@ -270,7 +276,7 @@ def set_pause_postcmd(pause_postcmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_resume_precmd(resume_precmd: Callable[[Any], str | None]) -> None: +def set_resume_precmd(resume_precmd: PrePostCmd) -> None: """ Set the function to call before the resume command. @@ -282,7 +288,7 @@ def set_resume_precmd(resume_precmd: Callable[[Any], str | None]) -> None: @usercommand @helparglist("") -def set_resume_postcmd(resume_postcmd: Callable[[Any], str | None]) -> None: +def set_resume_postcmd(resume_postcmd: PrePostCmd) -> None: """ Set the function to call after the resume command. @@ -400,7 +406,9 @@ def get_exp_data( @usercommand @helparglist("") @log_command_and_handle_exception -def get_spectrum_data(with_spec_zero: bool = True, with_time_bin_zero: bool = False) -> npt.NDArray: +def get_spectrum_data( + with_spec_zero: bool = True, with_time_bin_zero: bool = False +) -> npt.NDArray[np.float32]: """ Get the event mode spectrum data as ND array. diff --git a/src/genie_python/genie_pre_post_cmd_manager.py b/src/genie_python/genie_pre_post_cmd_manager.py index 8914167..0b7bc52 100644 --- a/src/genie_python/genie_pre_post_cmd_manager.py +++ b/src/genie_python/genie_pre_post_cmd_manager.py @@ -1,21 +1,45 @@ from builtins import object +from typing import Any class PrePostCmdManager(object): """ - A class to manager the precmd and postcmd commands such as used in begin, end, abort, resume, pause. + A class to manager the precmd and postcmd commands such as used in begin, end, abort, resume, + pause. """ - def __init__(self): - self.begin_precmd = lambda **pars: None - self.begin_postcmd = lambda **pars: None - self.abort_precmd = lambda **pars: None - self.abort_postcmd = lambda **pars: None - self.end_precmd = lambda **pars: None - self.end_postcmd = lambda **pars: None - self.pause_precmd = lambda **pars: None - self.pause_postcmd = lambda **pars: None - self.resume_precmd = lambda **pars: None - self.resume_postcmd = lambda **pars: None - self.cset_precmd = lambda **pars: True - self.cset_postcmd = lambda **pars: None + def begin_precmd(self, **kwargs: Any) -> str | None: + return + + def begin_postcmd(self, **kwargs: Any) -> str | None: + return + + def abort_precmd(self, **kwargs: Any) -> str | None: + return + + def abort_postcmd(self, **kwargs: Any) -> str | None: + return + + def end_precmd(self, **kwargs: Any) -> str | None: + return + + def end_postcmd(self, **kwargs: Any) -> str | None: + return + + def pause_precmd(self, **kwargs: Any) -> str | None: + return + + def pause_postcmd(self, **kwargs: Any) -> str | None: + return + + def resume_precmd(self, **kwargs: Any) -> str | None: + return + + def resume_postcmd(self, **kwargs: Any) -> str | None: + return + + def cset_precmd(self, **kwargs: Any) -> bool: + return True + + def cset_postcmd(self, **kwargs: Any) -> str | None: + return diff --git a/src/genie_python/genie_simulate_impl.py b/src/genie_python/genie_simulate_impl.py index 4262f98..5f9d8c6 100644 --- a/src/genie_python/genie_simulate_impl.py +++ b/src/genie_python/genie_simulate_impl.py @@ -1085,12 +1085,16 @@ def __init__( self.sample_pars = {} self.strict_block = strict_block self.logger = GenieLogger(sim_mode=True) + self.exp_data = None def set_instrument( self, pv_prefix: str, globs: dict | None, import_instrument_init: bool ) -> None: self.inst_prefix = pv_prefix + def get_instrument(self) -> str | None: + return self.inst_prefix + def prefix_pv_name(self, name: str) -> str: """Adds the instrument prefix to the specified PV""" if self.inst_prefix is not None and not name.startswith(self.inst_prefix): @@ -1100,7 +1104,7 @@ def prefix_pv_name(self, name: str) -> str: return name def set_pv_value( - self, name: str, value: str, wait: bool = False, is_local: bool = False + self, name: str, value: "PVValue", wait: bool = False, is_local: bool = False ) -> None: if is_local: name = self.prefix_pv_name(name) @@ -1119,7 +1123,7 @@ def get_pv_value( % (name, to_string, attempts, is_local) ) - def pv_exists(self, name: str) -> bool: + def pv_exists(self, name: str, is_local: bool = False) -> bool: return True def reload_current_config(self) -> None: @@ -1189,7 +1193,10 @@ def get_block_data(self, block: str, fail_fast: bool = False) -> dict(): ans["alarm"] = self.get_alarm_from_block(block) return ans - def set_multiple_blocks(self, names: list, values: list) -> None: + def get_pv_from_block(self, block_name: str) -> str: + return block_name + + def set_multiple_blocks(self, names: list[str], values: list["PVValue"]) -> None: temp = list(zip(names, values)) for name, value in temp: if name in self.block_dict: