From 13f74ef12d8e14345af76dbb77ff3cce37457f69 Mon Sep 17 00:00:00 2001
From: Daniel Stonier <d.stonier@gmail.com>
Date: Fri, 22 Nov 2019 20:11:49 -0500
Subject: [PATCH] [trees] play nicely, reset original signal handler after
 setup (#262)

---
 CHANGELOG.rst          |  2 +-
 py_trees/blackboard.py |  4 ----
 py_trees/trees.py      | 16 ++++++++++++----
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index e430c039..50570f85 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -3,7 +3,7 @@ Release Notes
 
 Forthcoming
 -----------
-* ...
+* [trees] play nicely, reset signal handlers after setup, `#262 <https://github.com/splintered-reality/py_trees/pull/262>`_
 
 2.0.1 (2019-11-19)
 ------------------
diff --git a/py_trees/blackboard.py b/py_trees/blackboard.py
index c656b14f..88e7372e 100755
--- a/py_trees/blackboard.py
+++ b/py_trees/blackboard.py
@@ -1122,12 +1122,8 @@ def unregister_key(self, key: str, clear: bool=True):
         Raises:
             KeyError if the key has not been previously registered
         """
-        print("DJS: Unregister Key")
-        print("DJS: Unregister Key: {}".format(key))
         key = Blackboard.absolute_name(super().__getattribute__("namespace"), key)
-        print("DJS: Unregister Key: {}".format(key))
         remapped_key = super().__getattribute__("remappings")[key]
-        print("DJS: Unregister Remapped Key: {}".format(remapped_key))
         super().__getattribute__("read").discard(key)  # doesn't throw exceptions if it not present
         super().__getattribute__("write").discard(key)
         super().__getattribute__("exclusive").discard(key)
diff --git a/py_trees/trees.py b/py_trees/trees.py
index 27a79573..54dabcc4 100644
--- a/py_trees/trees.py
+++ b/py_trees/trees.py
@@ -6,7 +6,6 @@
 ##############################################################################
 # Documentation
 ##############################################################################
-from distutils.core import setup
 
 """
 While a graph of connected behaviours and composites form a tree in their own right
@@ -26,6 +25,7 @@
 # Imports
 ##############################################################################
 
+import functools
 import os
 import signal
 import threading
@@ -75,8 +75,9 @@ def setup(root: behaviour.Behaviour,
     def on_timer_timed_out():
         os.kill(os.getpid(), _SIGNAL)
 
-    def signal_handler(unused_signum, unused_frame):
-        raise RuntimeError("tree setup timed out")
+    def signal_handler(unused_signum, unused_frame, original_signal_handler):
+        signal.signal(_SIGNAL, original_signal_handler)
+        raise RuntimeError("tree setup interrupted or timed out")
 
     def visited_setup():
         for node in root.iterate():
@@ -87,11 +88,18 @@ def visited_setup():
     if timeout == common.Duration.INFINITE:
         visited_setup()
     else:
-        signal.signal(_SIGNAL, signal_handler)
+        original_signal_handler = signal.getsignal(_SIGNAL)
+        signal.signal(
+            _SIGNAL,
+            functools.partial(
+                signal_handler,
+                original_signal_handler=original_signal_handler)
+        )
         timer = threading.Timer(interval=timeout, function=on_timer_timed_out)
         timer.start()
         visited_setup()
         timer.cancel()  # this only works if the timer is still waiting
+        signal.signal(_SIGNAL, original_signal_handler)
 
 ##############################################################################
 # Trees