Merge pull request #4 from spc-group/weakref

Added an argument *keep_references* to Registry() that forces weak re…

README.md

@@ -223,6 +223,39 @@ motor.set(5).wait()
 
 ```
 
+Keeping References
+------------------
+
+It may be useful to not actually keep a strong reference to the
+``OphydObject``s. This means that if all other references to the
+object are removed, the device may be dropped from the registry.
+
+By default, the registry keeps direct references to the objects that
+get registered, but if initalized with ``keep_references=False`` the
+Registry will not keep these references. Instead, **it is up to you to
+keep references to the registered objects**.
+
+```python
+
+# Create two registers with both referencing behaviors
+ref_registry = Registry(keep_references=True)
+noref_registry = Registry(keep_references=False)
+motor = EpicsMotor(...)
+
+# Check if we can get the motor (should be no problem)
+ref_registry[motor.name]  # <- succeeds
+noref_registry[motor.name]  # <- succeeds
+
+# Delete the object and try again
+del motor
+gc.collect()  # <- make sure it's *really* gone
+
+# Check again if we can get the motor (now it gets fun)
+ref_registry[motor.name]  # <- succeeds
+noref_registry[motor.name]  # <- raises ComponentNotFound
+
+```
+
 Integrating with Typhos
 -----------------------
 

src/ophydregistry/registry.py

@@ -3,6 +3,7 @@
 from collections import OrderedDict
 from itertools import chain
 from typing import List, Mapping, Optional
+from weakref import WeakValueDictionary, WeakSet
 
 from ophyd import ophydobj
 
@@ -103,23 +104,31 @@ class Registry:
     use_typhos
       If true, items added to this registry will also be added to the
       Typhos registry for inclusion in PyDM windows.
+    keep_references
+      If false, items will be dropped from this registry if the only
+      reference comes from this registry. Relies on the garbage
+      collector, so to force cleanup use ``gc.collect()``.
 
     """
 
-    use_typhos: bool = False
+    use_typhos: bool
+    keep_references: bool
     _auto_register: bool
 
     # components: Sequence
     _objects_by_name: Mapping
     _objects_by_label: Mapping
 
-    def __init__(self, auto_register: bool = True, use_typhos: bool = False):
+    def __init__(
+        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:
             raise ModuleNotFoundError("No module named 'typhos'")
         # Set up empty lists and things for registering components
-        self.clear()
+        self.keep_references = keep_references
         self.use_typhos = use_typhos
+        self.clear()
         self.auto_register = auto_register
 
     @property
@@ -202,8 +211,11 @@ def clear(self, clear_typhos: bool = True):
           *self.use_typhos* is false.
 
         """
-        self._objects_by_name = OrderedDict()
         self._objects_by_label = OrderedDict()
+        if self.keep_references:
+            self._objects_by_name = OrderedDict()
+        else:
+            self._objects_by_name = WeakValueDictionary()
         if clear_typhos and self.use_typhos:
             typhos.plugins.core.signal_registry.clear()
 
@@ -496,7 +508,10 @@ def register(self, component: ophydobj.OphydObject) -> ophydobj.OphydObject:
             # Create a set for this device's labels if it doesn't exist
             for label in getattr(component, "_ophyd_labels_", []):
                 if label not in self._objects_by_label.keys():
-                    self._objects_by_label[label] = set()
+                    if self.keep_references:
+                        self._objects_by_label[label] = set()
+                    else:
+                        self._objects_by_label[label] = WeakSet()
                 self._objects_by_label[label].add(component)
             # Register this object with Typhos
             if self.use_typhos:

src/ophydregistry/tests/test_instrument_registry.py

@@ -429,3 +429,21 @@ def test_pop_default(registry):
         registry[motor.name]
     with pytest.raises(ComponentNotFound):
         registry["motors"]
+
+
+def test_weak_references():
+    """Check that we can make a registry that automatically drops
+    objects that are only referenced by this registry.
+
+    """
+    motor = sim.SynAxis(name="weak_motor", labels={"motors"})
+    registry = Registry(keep_references=False)
+    registry.register(motor)
+    # Can we still find the object if the test has a reference?
+    assert registry.find("weak_motor") is motor
+    # Delete the original object
+    del motor
+    gc.collect()
+    # Check that it's not in the registry anymore
+    with pytest.raises(ComponentNotFound):
+        registry.find("weak_motor")