Add {add,remove}_instrument, deprecate current_instruments

For gh-257

docs/source/reference-core.rst

@@ -1702,8 +1702,14 @@ handling exceptions raised by instruments is to (a) dump a stack trace
 to stderr and (b) disable the offending instrument.
 
 You can register an initial list of instruments by passing them to
-:func:`trio.run`. :func:`current_instruments` lets you introspect and
-modify this list at runtime from inside trio:
+:func:`trio.run`. :func:`add_instrument` and
+:func:`remove_instrument` let you add and remove instruments at
+runtime. There's also :func:`current_instruments`, which is deprecated
+and will be removed soon.
+
+.. autofunction:: add_instrument
+
+.. autofunction:: remove_instrument
 
 .. autofunction:: current_instruments
 

trio/_core/_run.py

@@ -16,6 +16,7 @@
 from async_generator import async_generator, yield_
 
 from .._util import acontextmanager
+from .._deprecate import deprecated
 
 from .. import _core
 from ._exceptions import (
@@ -1299,6 +1300,41 @@ def instrument(self, method_name, *args):
                 sys.stderr.write("Instrument has been disabled.\n")
 
     @_public
+    def add_instrument(self, instrument):
+        """Start instrumenting the current run loop with the given instrument.
+
+        Args:
+          instrument (trio.abc.Instrument): The instrument to activate.
+
+        If ``instrument`` is already active, does nothing.
+
+        """
+        if instrument not in self.instruments:
+            self.instruments.append(instrument)
+
+    @_public
+    def remove_instrument(self, instrument):
+        """Stop instrumenting the current run loop with the given instrument.
+
+        Args:
+          instrument (trio.abc.Instrument): The instrument to de-activate.
+
+        Raises:
+          KeyError: if the instrument is not currently active. This could
+              occur either because you never added it, or because you added it
+              and then it raised an unhandled exception and was automatically
+              deactivated.
+
+        """
+        # We're moving 'instruments' to being a set, so raise KeyError like
+        # set.remove does.
+        try:
+            self.instruments.remove(instrument)
+        except ValueError as exc:
+            raise KeyError(*exc.args)
+
+    @_public
+    @deprecated("0.2.0", issue=257, instead="{add,remove}_instrument")
     def current_instruments(self):
         """Returns the list of currently active instruments.
 

trio/_core/tests/test_run.py

@@ -434,7 +434,7 @@ def filter_tasks(self, tasks):
                 yield item
 
 
-def test_instruments():
+def test_instruments_deprecated(recwarn):
     r1 = TaskRecorder()
     r2 = TaskRecorder()
     r3 = TaskRecorder()
@@ -465,6 +465,40 @@ async def main():
     assert list(r1.filter_tasks([task])) == expected
 
 
+def test_instruments(recwarn):
+    r1 = TaskRecorder()
+    r2 = TaskRecorder()
+    r3 = TaskRecorder()
+
+    async def main():
+        for _ in range(4):
+            await _core.yield_briefly()
+        # replace r2 with r3, to test that we can manipulate them as we go
+        _core.remove_instrument(r2)
+        with pytest.raises(KeyError):
+            _core.remove_instrument(r2)
+        # add is idempotent
+        _core.add_instrument(r3)
+        _core.add_instrument(r3)
+        for _ in range(1):
+            await _core.yield_briefly()
+        return _core.current_task()
+
+    task = _core.run(main, instruments=[r1, r2])
+    # It sleeps 5 times, so it runs 6 times
+    expected = (
+        [("before_run",)] +
+        6 * [("schedule", task),
+             ("before", task),
+             ("after", task)] + [("after_run",)]
+    )
+    assert len(r1.record) > len(r2.record) > len(r3.record)
+    assert r1.record == r2.record + r3.record
+    # Need to filter b/c there's also the system task bumping around in the
+    # record:
+    assert list(r1.filter_tasks([task])) == expected
+
+
 def test_instruments_interleave():
     tasks = {}