Source code for xrd_tools.device_manager

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Callable

from . import device_factory, utils
from .config import DEVICES, DEVICE_TYPES, ENCODING, JSON_INDENT
from .device import HTChamber, LabDiffractometer

logger = logging.getLogger(__name__)


[docs]@dataclass class DeviceManager: """Manage Devices.""" file_path: str = DEVICES device_types: dict[str, str] = field(default_factory=lambda: DEVICE_TYPES) encoding: str = ENCODING indent: str | None = JSON_INDENT separator: str = ": " def __post_init__(self): # Register device types self._register_device_types() # Load and initialise devices dev_tuples = [device_factory.create(item) for item in self._load_data()] self.device_dict = {item[1]: item[0] for item in dev_tuples} def _validate_device_type(self, device_type: str) -> None: """Validates a cpecified device type. Args: device_type (str): Device type to be validated. Raises: ValueError: If the device type is invalid. """ if device_type not in device_factory.device_creation_funcs.keys(): raise ValueError(f"Invalid device_type provided {device_type!r}.") def _register_device_types(self) -> None: """Register device types.""" device_factory.register(DEVICE_TYPES["lab_xrd"], LabDiffractometer) device_factory.register(DEVICE_TYPES["ht_chamber"], HTChamber) def _load_data(self) -> list[dict[str, any]]: """Read device dictionaries from json file.""" if not os.path.isfile(self.file_path): logger.debug("No devices registered.") return [] with open(self.file_path, "r", encoding=self.encoding) as fobj: json_str = fobj.read() logger.debug(f"Loaded devices from '{os.path.abspath(self.file_path)}'") return json.loads(json_str) def _write_data(self, file_path: str = None) -> None: """Write json string of device_dict to file.""" # Ensure file is provided and parent directory exists if file_path is None: file_path = self.file_path utils.make_dirs(file_path) # Retrieve data, add device_type and device_id in addition to the class dict data = [] for device_id, dev_obj in self.device_dict.items(): kwargs = {"device_type": dev_obj.device_type, "device_id": device_id} kwargs.update(dev_obj.__dict__.copy()) # Drop None values drop = [k for k, v in kwargs.items() if v is None] for k in drop: kwargs.pop(k) data.append(kwargs) # Create JSON string and write to file taking indent into account utils.write_to_json(file_path, data, indent=self.indent, encoding=self.encoding)
[docs] def add_device( self, device_id: str, dev_obj: device_factory.Device, to_file: bool = True ) -> None: """Add device kwargs to json file.""" if device_id is None: raise ValueError("Valid device ID must be provided.") elif device_id in self.get_id_list(): raise ValueError(f"Device with ID {dev_obj.device_id!r} already existing.") self.device_dict[device_id] = dev_obj logger.info(f"Added device: {device_id!r} ({dev_obj.device_type})") if to_file: self._write_data()
[docs] def get_creation_function( self, device_type: str ) -> Callable[..., device_factory.Device]: """Get creation function for device type.""" return device_factory.device_creation_funcs[device_type]
[docs] def get_device(self, device_id) -> device_factory.Device: """Return device object for provided ID.""" return self.device_dict[device_id]
[docs] def get_id_list(self, device_type: str = None): """Return list of known device IDs. Filters for specific device type if one is provided. """ if device_type is not None: return [ k for k, v in self.device_dict.items() if v.device_type == device_type ] return self.device_dict.keys()
[docs] def get_type_device_list(self): """Returns list of known device IDs and their type.""" return [ f"{d.device_type}{self.separator}{k}" for k, d in self.device_dict.items() ]
[docs] def get_type(self, device_id: str) -> str: """Get the type of a registered device. Args: device_id (str): The ID of a registered device whose type is returned. Returns: str: The device type of the device with specified ID. """ return self.device_dict[device_id].device_type
[docs] def get_device_ids(self, device_type: str = None) -> list[str]: """Get a list of device IDs, optionally of a certain type. Args: device_type (str): The type of device IDs to be returned. Returns: str: List of all device IDs if no device_type is specified, otherwise a list with all devices of the corresponding type. """ if device_type is None: return self.device_dict.keys() devices = self.get_devices(device_type=device_type) return [k for k, v in self.device_dict.items() if v in devices]
[docs] def get_devices(self, device_type: str = None) -> list[device_factory.Device]: """Get a list of devices, optionally of a certain type. Args: device_type (str): The type of devices to be returned. Returns: str: List of all devices if no device_type is specified, otherwise a list with all devices of the corresponding type. """ if device_type is None: return self.device_dict.values() self._validate_device_type(device_type) creation_func = self.get_creation_function(device_type) return [d for d in self.device_dict.values() if isinstance(d, creation_func)]
[docs] def remove_device(self, device_id: str, to_file: bool = True) -> None: """Remove a registered device (and write updated devices to file). Args: device_id (str): The ID of the device to be removed. to_file (bool): Flag to indicate to write the new devices to the JSON file. """ device_type = self.device_dict[device_id].device_type self.device_dict.pop(device_id) logger.info(f"Removed device: {device_id!r} ({device_type})") if to_file: self._write_data()