Merge pull request #8312 from jenshnielsen/close_subclass

Instrument Add option to only close specific subclass

Author: astafan8

docs/changes/newsfragments/8312.improved

@@ -0,0 +1,4 @@
+:meth:`.Instrument.close_all` gained an ``only_subclasses`` keyword argument. When
+set to ``True``, only instruments that are instances of the class (or its subclasses)
+on which ``close_all`` is called are closed, leaving other registered instruments open.
+The ``log_status`` and ``only_subclasses`` arguments are now keyword-only.

src/qcodes/instrument/instrument.py

@@ -182,7 +182,9 @@ def close(self) -> None:
     @classmethod
     def close_all(
         cls,
+        *,
         log_status: bool = False,
+        only_subclasses: bool = False,
     ) -> None:
         """
         Try to close all instruments registered in
@@ -196,16 +198,25 @@ def close_all(
         Args:
             log_status: If True, log the status of closing each instrument. Set this to False
               if you want to avoid logging during interpreter shutdown, which can cause errors.
+            only_subclasses: If True, only close instruments that are subclasses of the class
+              on which this method is called. If False, close all instruments regardless of class.
 
         """
         if log_status:
             log.info("Closing all registered instruments")
         for inststr in list(cls._all_instruments):
             try:
                 inst: Instrument = cls.find_instrument(inststr)
-                if log_status:
-                    log.info("Closing %s", inststr)
-                inst.close()
+                if only_subclasses and issubclass(type(inst), cls):
+                    should_be_closed = True
+                elif not only_subclasses:
+                    should_be_closed = True
+                else:
+                    should_be_closed = False
+                if should_be_closed:
+                    if log_status:
+                        log.info("Closing %s", inststr)
+                    inst.close()
             except Exception:
                 if log_status:
                     log.exception("Failed to close %s, ignored", inststr)

tests/test_instrument.py

@@ -7,6 +7,7 @@
 import contextlib
 import gc
 import io
+import logging
 import re
 import weakref
 from typing import TYPE_CHECKING, Any, assert_type
@@ -22,6 +23,7 @@
     find_or_create_instrument,
 )
 from qcodes.instrument_drivers.mock_instruments import (
+    DummyBase,
     DummyChannelInstrument,
     DummyFailingInstrument,
     DummyInstrument,
@@ -450,6 +452,101 @@ def test_recreate(request: FixtureRequest) -> None:
     assert instr not in Instrument._all_instruments.values()
 
 
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_closes_all_instruments() -> None:
+    """``close_all`` closes every registered instrument by default."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+    parabola = MockParabola("parabola")
+
+    assert Instrument.is_valid(dummy)
+    assert Instrument.is_valid(parabola)
+
+    Instrument.close_all()
+
+    assert not Instrument.is_valid(dummy)
+    assert not Instrument.is_valid(parabola)
+    assert Instrument._all_instruments == WeakValueDictionary()
+
+
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_only_subclasses_from_leaf_class() -> None:
+    """``only_subclasses`` on a leaf class leaves sibling classes open."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+    parabola = MockParabola("parabola")
+
+    # DummyInstrument and MockParabola are siblings (both subclass DummyBase),
+    # so closing only DummyInstrument subclasses must leave the parabola open.
+    DummyInstrument.close_all(only_subclasses=True)
+
+    assert not Instrument.is_valid(dummy)
+    assert Instrument.is_valid(parabola)
+
+    # The remaining instrument can still be closed with a plain close_all.
+    Instrument.close_all()
+    assert not Instrument.is_valid(parabola)
+    assert Instrument._all_instruments == WeakValueDictionary()
+
+
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_only_subclasses_from_base_class() -> None:
+    """``only_subclasses`` closes instances of the class and its subclasses."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+    parabola = MockParabola("parabola")
+
+    # Both DummyInstrument and MockParabola are subclasses of DummyBase, so both
+    # are closed when calling close_all on the shared base class.
+    DummyBase.close_all(only_subclasses=True)
+
+    assert not Instrument.is_valid(dummy)
+    assert not Instrument.is_valid(parabola)
+    assert Instrument._all_instruments == WeakValueDictionary()
+
+
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_only_subclasses_false_closes_everything() -> None:
+    """``only_subclasses=False`` closes all instruments regardless of class."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+    parabola = MockParabola("parabola")
+
+    DummyInstrument.close_all(only_subclasses=False)
+
+    assert not Instrument.is_valid(dummy)
+    assert not Instrument.is_valid(parabola)
+    assert Instrument._all_instruments == WeakValueDictionary()
+
+
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_log_status(caplog: pytest.LogCaptureFixture) -> None:
+    """``log_status=True`` logs the closing of each instrument."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+
+    with caplog.at_level(logging.INFO, logger="qcodes.instrument.instrument"):
+        Instrument.close_all(log_status=True)
+
+    assert "Closing all registered instruments" in caplog.text
+    assert "Closing dummy" in caplog.text
+    assert not Instrument.is_valid(dummy)
+
+
+@pytest.mark.usefixtures("close_before_and_after")
+def test_close_all_no_log_by_default(caplog: pytest.LogCaptureFixture) -> None:
+    """``close_all`` does not log anything when ``log_status`` is not set."""
+    dummy = DummyInstrument(name="dummy", gates=["dac1"])
+
+    with caplog.at_level(logging.INFO, logger="qcodes.instrument.instrument"):
+        Instrument.close_all()
+
+    assert "Closing all registered instruments" not in caplog.text
+    assert "Closing dummy" not in caplog.text
+    assert not Instrument.is_valid(dummy)
+
+
+def test_close_all_only_accepts_keyword_arguments() -> None:
+    """The ``close_all`` options are keyword-only."""
+    with pytest.raises(TypeError):
+        Instrument.close_all(True)  # type: ignore[misc]
+
+
 def test_instrument_metadata(request: FixtureRequest) -> None:
     metadatadict = {1: "data", "some": "data"}
     instrument = DummyInstrument(