From d7f39bd893fd163909cc448d015344ecdb8428ca Mon Sep 17 00:00:00 2001 From: Bart Schilperoort Date: Thu, 11 Jan 2024 09:23:23 +0100 Subject: [PATCH] Fix typing and formatting issues --- PyStemmusScope/bmi/docker_process.py | 43 ++++++++++++++-------------- PyStemmusScope/bmi/implementation.py | 36 ++++++++++++++++------- PyStemmusScope/bmi/local_process.py | 11 ++++--- PyStemmusScope/soil_io.py | 2 +- 4 files changed, 53 insertions(+), 39 deletions(-) diff --git a/PyStemmusScope/bmi/docker_process.py b/PyStemmusScope/bmi/docker_process.py index 2683c671..1d8cbaf6 100644 --- a/PyStemmusScope/bmi/docker_process.py +++ b/PyStemmusScope/bmi/docker_process.py @@ -1,10 +1,11 @@ """The Docker STEMMUS_SCOPE model process wrapper.""" +import os +import warnings +from pathlib import Path from time import sleep from typing import Any -import warnings from PyStemmusScope.config_io import read_config -from pathlib import Path -import os + try: import docker @@ -30,7 +31,9 @@ def make_docker_vols_binds(cfg_file: str) -> tuple[list[str], list[str]]: if not cfg_dir.is_relative_to(cfg["InputPath"]): volumes.append(str(cfg_dir)) binds.append(f"{str(cfg_dir)}:{str(cfg_dir)}") - if (not Path(cfg["InputPath"]).is_relative_to(cfg_dir)) or (Path(cfg["InputPath"]) == cfg_dir): + if (not Path(cfg["InputPath"]).is_relative_to(cfg_dir)) or ( + Path(cfg["InputPath"]) == cfg_dir + ): volumes.append(cfg["InputPath"]) binds.append(f"{cfg['InputPath']}:{cfg['InputPath']}") if not Path(cfg["OutputPath"]).is_relative_to(cfg_dir): @@ -80,8 +83,9 @@ def wait_for_model(phrase: bytes, socket: Any) -> None: class StemmusScopeDocker: """Communicate with a STEMMUS_SCOPE Docker container.""" + # Default image, can be overridden with config: - compatible_tags = ("1.5.0", ) + compatible_tags = ("1.5.0",) _process_ready_phrase = b"Select BMI mode:" _process_finalized_phrase = b"Finished clean up." @@ -104,53 +108,50 @@ def __init__(self, cfg_file: str): detach=True, user=os.getuid(), # ensure correct user for writing files. volumes=vols, - host_config=self.client.create_host_config(binds=binds) + host_config=self.client.create_host_config(binds=binds), ) self.running = False - - def wait_for_model(self): + + def wait_for_model(self) -> None: """Wait for the model to be ready to receive (more) commands.""" wait_for_model(self._process_ready_phrase, self.socket) - - def is_alive(self): + + def is_alive(self) -> bool: """Return if the process is alive.""" return self.running - - def initialize(self): + + def initialize(self) -> None: """Initialize the model and wait for it to be ready.""" if self.is_alive(): self.client.stop(self.container_id) self.client.start(self.container_id) self.socket = self.client.attach_socket( - self.container_id, {'stdin': 1, 'stdout': 1, 'stream':1} + self.container_id, {"stdin": 1, "stdout": 1, "stream": 1} ) self.wait_for_model() os.write( self.socket.fileno(), - bytes(f'initialize "{self.cfg_file}"\n', encoding="utf-8") + bytes(f'initialize "{self.cfg_file}"\n', encoding="utf-8"), ) self.wait_for_model() self.running = True - def update(self): + def update(self) -> None: """Update the model and wait for it to be ready.""" if self.is_alive(): - os.write( - self.socket.fileno(), - b'update\n' - ) + os.write(self.socket.fileno(), b"update\n") self.wait_for_model() else: msg = "Docker container is not alive. Please restart the model." raise ConnectionError(msg) - def finalize(self): + def finalize(self) -> None: """Finalize the model.""" if self.is_alive(): - os.write(self.socket.fileno(),b'finalize\n') + os.write(self.socket.fileno(), b"finalize\n") wait_for_model(self._process_finalized_phrase, self.socket) sleep(0.5) # Ensure the container can stop cleanly. self.client.stop(self.container_id) diff --git a/PyStemmusScope/bmi/implementation.py b/PyStemmusScope/bmi/implementation.py index 2b8b6b03..b8c8934f 100644 --- a/PyStemmusScope/bmi/implementation.py +++ b/PyStemmusScope/bmi/implementation.py @@ -2,7 +2,8 @@ import os import sys from pathlib import Path -from typing import Literal, Protocol +from typing import Literal +from typing import Protocol from typing import Union import h5py import numpy as np @@ -124,7 +125,7 @@ def set_variable( return state -def get_run_mode(config: dict) -> str: +def get_run_mode(config: dict) -> Literal["exe", "docker"]: """Get the run mode (docker or EXE) from the config file. Args: @@ -150,6 +151,7 @@ def get_run_mode(config: dict) -> str: class StemmusScopeProcess(Protocol): """Protocol for communicating with the model process.""" + def __init__(self, cfg_file: str) -> None: """Initialize the process class (e.g. create the container).""" ... @@ -171,11 +173,13 @@ def finalize(self) -> None: ... -def load_process(mode: Literal["exe", "docker"]) -> type[StemmusScopeProcess]: - """Load the right STEMMUS_SCOPE process.""" +def start_process(mode: Literal["exe", "docker"], cfg_file: str) -> StemmusScopeProcess: + """Start the right STEMMUS_SCOPE process.""" if mode == "docker": try: - from PyStemmusScope.bmi.docker_process import StemmusScopeDocker as Process + from PyStemmusScope.bmi.docker_process import StemmusScopeDocker + + return StemmusScopeDocker(cfg_file=cfg_file) except ImportError as err: msg = ( "The docker python package is not available." @@ -183,11 +187,12 @@ def load_process(mode: Literal["exe", "docker"]) -> type[StemmusScopeProcess]: ) raise ImportError(msg) from err elif mode == "exe": - from PyStemmusScope.bmi.local_process import LocalStemmusScope as Process + from PyStemmusScope.bmi.local_process import LocalStemmusScope + + return LocalStemmusScope(cfg_file=cfg_file) else: msg = "Unknown mode." raise ValueError(msg) - return Process class StemmusScopeBmi(InapplicableBmiMethods, Bmi): @@ -199,7 +204,7 @@ class StemmusScopeBmi(InapplicableBmiMethods, Bmi): state_file: Union[Path, None] = None _run_mode: Union[str, None] = None - _process: Union[type[StemmusScopeProcess], None] = None + _process: Union[StemmusScopeProcess, None] = None def initialize(self, config_file: str) -> None: """Perform startup tasks for the model. @@ -213,7 +218,7 @@ def initialize(self, config_file: str) -> None: self._run_mode = get_run_mode(self.config) self.state_file = Path(self.config["OutputPath"]) / "STEMMUS_SCOPE_state.mat" - self._process = load_process(self._run_mode)(cfg_file=config_file) + self._process = start_process(self._run_mode, config_file) self._process.initialize() def update(self) -> None: @@ -221,7 +226,12 @@ def update(self) -> None: if self.state is not None: self.state = self.state.close() # Close file to allow matlab to write - self._process.update() + if self._process is not None: + self._process.update() + else: + msg = "The STEMMUS_SCOPE process is not running/connected. Can't update!" + raise ValueError(msg) + self.state = load_state(self.config) def update_until(self, time: float) -> None: @@ -235,7 +245,11 @@ def update_until(self, time: float) -> None: def finalize(self) -> None: """Finalize the STEMMUS_SCOPE model.""" - self._process.finalize() + if self._process is not None: + self._process.finalize() + else: + msg = "The STEMMUS_SCOPE process is not running/connected. Can't finalize!" + raise ValueError(msg) def get_component_name(self) -> str: """Name of the component. diff --git a/PyStemmusScope/bmi/local_process.py b/PyStemmusScope/bmi/local_process.py index b65cab9b..50c805da 100644 --- a/PyStemmusScope/bmi/local_process.py +++ b/PyStemmusScope/bmi/local_process.py @@ -1,9 +1,9 @@ """The local STEMMUS_SCOPE model process wrapper.""" -from pathlib import Path +import os import subprocess +from pathlib import Path from typing import Union from PyStemmusScope.config_io import read_config -import os def is_alive(process: Union[subprocess.Popen, None]) -> subprocess.Popen: @@ -42,6 +42,7 @@ def find_exe(config: dict) -> str: class LocalStemmusScope: """Communicate with the local STEMMUS_SCOPE executable file.""" + def __init__(self, cfg_file: str) -> None: """Initialize the process.""" self.cfg_file = cfg_file @@ -60,7 +61,7 @@ def __init__(self, cfg_file: str) -> None: ) wait_for_model(self.matlab_process) - + def is_alive(self) -> bool: """Return if the process is alive.""" try: @@ -74,11 +75,10 @@ def initialize(self) -> None: self.matlab_process = is_alive(self.matlab_process) self.matlab_process.stdin.write( # type: ignore - bytes(f'initialize "{self.cfg_file}"\n', encoding="utf-8") + bytes(f'initialize "{self.cfg_file}"\n', encoding="utf-8") ) wait_for_model(self.matlab_process) - def update(self) -> None: """Update the model and wait for it to be ready.""" if self.matlab_process is None: @@ -89,7 +89,6 @@ def update(self) -> None: self.matlab_process.stdin.write(b"update\n") # type: ignore wait_for_model(self.matlab_process) - def finalize(self) -> None: """Finalize the model.""" self.matlab_process = is_alive(self.matlab_process) diff --git a/PyStemmusScope/soil_io.py b/PyStemmusScope/soil_io.py index 37ea9558..cc3d245d 100644 --- a/PyStemmusScope/soil_io.py +++ b/PyStemmusScope/soil_io.py @@ -1,6 +1,6 @@ """Module for the soil data IO of PyStemmusScope.""" -from pathlib import Path from collections.abc import Iterable +from pathlib import Path import hdf5storage import numpy as np import xarray as xr