diff --git a/README.md b/README.md index 8c95a51..93406f0 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,15 @@ method: registry.clear() ``` +To **remove disconnected devices** from the registry, use the ``pop_disconnected()`` method with an optional timeout: + +```python + +# Wait 5 seconds to give devices a chance to connect +disconnected_devices = registry.pop_disconnected(timeout=5) + +``` + To **remove individual objects**, use either the *del* keyword, or the ``pop()`` method. These approaches work with either the ``OphydObject`` instance itself, or the instance's name: diff --git a/pyproject.toml b/pyproject.toml index 1c713f3..f8c29b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ophyd-registry" -version = "1.1.0" +version = "1.2.0" authors = [ { name="Mark Wolfman", email="wolfman@anl.gov" }, ] @@ -19,7 +19,7 @@ dependencies = ["ophyd"] [project.optional-dependencies] -dev = ["black", "isort", "pytest", "build", "twine", "flake8", "ruff"] +dev = ["black", "isort", "pytest", "build", "twine", "flake8", "ruff", "pytest-mock"] [project.urls] "Homepage" = "https://github.com/spc-group/ophyd-registry" diff --git a/src/ophydregistry/registry.py b/src/ophydregistry/registry.py index 520a0f8..7ba1e0f 100644 --- a/src/ophydregistry/registry.py +++ b/src/ophydregistry/registry.py @@ -1,9 +1,10 @@ import logging +import time import warnings from collections import OrderedDict from itertools import chain -from typing import List, Mapping, Optional -from weakref import WeakValueDictionary, WeakSet +from typing import List, Mapping, Optional, Set +from weakref import WeakSet, WeakValueDictionary from ophyd import ophydobj @@ -114,13 +115,17 @@ class Registry: use_typhos: bool keep_references: bool _auto_register: bool + _valid_classes: Set[type] = {ophydobj.OphydObject} # components: Sequence _objects_by_name: Mapping _objects_by_label: Mapping def __init__( - self, auto_register: bool = True, use_typhos: bool = False, keep_references: bool = True + self, + auto_register: bool = True, + use_typhos: bool = False, + keep_references: bool = True, ): # Check that Typhos is installed if needed if use_typhos and not typhos_available: @@ -219,6 +224,35 @@ def clear(self, clear_typhos: bool = True): if clear_typhos and self.use_typhos: typhos.plugins.core.signal_registry.clear() + def pop_disconnected(self, timeout: float = 0.0) -> List: + """Remove any registered objects that are disconnected. + + Parameters + ========== + timeout + How long to wait for devices to connect, in seconds. + + Returns + ======= + disconnected + The root level devices that were removed. + + """ + remaining = [dev for dev in self.root_devices] + t0 = time.monotonic() + timeout_reached = False + while not timeout_reached: + # Remove any connected devices for the running list + remaining = [dev for dev in remaining if not dev.connected] + if len(remaining) == 0: + # All devices are connected, so just end early. + break + time.sleep(min((0.05, timeout / 10.0))) + timeout_reached = (time.monotonic() - t0) > timeout + # Remove unconnected devices from the registry + popped = [self.pop(dev) for dev in remaining] + return popped + @property def component_names(self): return set(self._objects_by_name.keys()) @@ -304,9 +338,22 @@ def find( result = None return result + def _is_resolved(self, obj): + """Is the object already resolved into an ophyd device, etc. + + This method checks the type of the object. To extend this to + other types of objects, override this objects + ``_valid_classes`` attribute with a new set. + + """ + for cls in self._valid_classes: + if isinstance(obj, cls): + return True + return False + def _findall_by_label(self, label, allow_none): # Check for already created ophyd objects (return as is) - if isinstance(label, ophydobj.OphydObject): + if self._is_resolved(label): yield label return # Recursively get lists of components @@ -333,7 +380,7 @@ def _findall_by_label(self, label, allow_none): def _findall_by_name(self, name): # Check for already created ophyd objects (return as is) - if isinstance(name, ophydobj.OphydObject): + if self._is_resolved(name): yield name return # Check for an edge case with EpicsMotor objects (user_readback name is same as parent) diff --git a/src/ophydregistry/tests/test_instrument_registry.py b/src/ophydregistry/tests/test_instrument_registry.py index 281ed49..0a3003d 100644 --- a/src/ophydregistry/tests/test_instrument_registry.py +++ b/src/ophydregistry/tests/test_instrument_registry.py @@ -1,5 +1,7 @@ -import gc import logging +import time +from concurrent.futures import ThreadPoolExecutor +from unittest import mock import pytest from ophyd import Device, EpicsMotor, sim @@ -10,6 +12,7 @@ @pytest.fixture() def registry(): reg = Registry(auto_register=False, use_typhos=False) + reg._valid_classes = {mock.MagicMock, *reg._valid_classes} try: yield reg finally: @@ -447,3 +450,53 @@ def test_weak_references(): # Check that it's not in the registry anymore with pytest.raises(ComponentNotFound): registry.find("weak_motor") + + +@pytest.fixture() +def motors(mocker): + mocker.patch("ophyd.epics_motor.EpicsMotor.connected", new=True) + good_motor = EpicsMotor(name="good_motor") + good_motor.connected = True + bad_motor = EpicsMotor(name="bad_motor") + bad_motor.connected = False + return (good_motor, bad_motor) + + +def test_pop_disconnected(registry, motors): + """Check that we can remove disconnected devices.""" + good_motor, bad_motor = motors + registry.register(good_motor) + registry.register(bad_motor) + # Check that the disconnected device gets removed + popped = registry.pop_disconnected() + with pytest.raises(ComponentNotFound): + registry["bad_motor"] + # Check that the popped device was returned + assert len(popped) == 1 + assert popped[0] is bad_motor + # Check that the connected device is still in the registry + assert registry["good_motor"] is good_motor + + +def test_pop_disconnected_with_timeout(registry, motors): + """Check that we can apply a timeout when waiting for disconnected + devices. + + """ + good_motor, bad_motor = motors + good_motor.connected = False # It starts disconnected + # Register the devices + registry.register(good_motor) + registry.register(bad_motor) + + # Remove the devices with a timeout + def make_connected(dev, wait): + time.sleep(wait) + dev.connected = True + + with ThreadPoolExecutor(max_workers=1) as executor: + # Make the connection happen after 50 ms + executor.submit(make_connected, good_motor, 0.15) + registry.pop_disconnected(timeout=0.3) + # Check that the connected device is still in the registry + assert registry["good_motor"] is good_motor