From f5d22a3ca19bd8cfb2aec09aa2d1f0c2d5162149 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 29 Apr 2018 19:41:23 +0100
Subject: [PATCH] core: support deleting handlers, make Receiver.close()
 unregister

---
 ansible_mitogen/target.py | 2 +-
 docs/api.rst              | 7 +++++++
 mitogen/core.py           | 6 ++++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/ansible_mitogen/target.py b/ansible_mitogen/target.py
index 3020595e..b42ecf31 100644
--- a/ansible_mitogen/target.py
+++ b/ansible_mitogen/target.py
@@ -105,7 +105,7 @@ def _get_file(context, path, out_fp):
             kwargs={
                 'size': len(s),
             }
-        )
+        ).close()
         out_fp.write(s)
 
     ok = out_fp.tell() == metadata['size']
diff --git a/docs/api.rst b/docs/api.rst
index 17b1b38d..923c368d 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -470,6 +470,13 @@ Router Class
         :return:
             `handle`, or if `handle` was ``None``, the newly allocated handle.
 
+    .. method:: del_handler (handle)
+
+        Remove the handle registered for `handle`
+
+        :raises KeyError:
+            The handle wasn't registered.
+
     .. method:: _async_route(msg, stream=None)
 
         Arrange for `msg` to be forwarded towards its destination. If its
diff --git a/mitogen/core.py b/mitogen/core.py
index 2e43ed65..d87baf36 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -436,6 +436,9 @@ class Receiver(object):
             self.notify(self)
 
     def close(self):
+        if self.handle:
+            self.router.del_handler(self.handle)
+            self.handle = None
         self._latch.put(Message.dead())
 
     def empty(self):
@@ -1273,6 +1276,9 @@ class Router(object):
         self.broker.start_receive(stream)
         listen(stream, 'disconnect', lambda: self.on_stream_disconnect(stream))
 
+    def del_handler(self, handle):
+        del self._handle_map[handle]
+
     def add_handler(self, fn, handle=None, persist=True,
                     policy=None, respondent=None):
         handle = handle or self._last_handle.next()