Add OptManager.subscribe, use it to clean up palettes in console

.subscribe lets you subscribe a function to a specified set of options.

4 files changed, 61 insertions, 22 deletions

diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py
index 56122baa..3ce76a45 100644
--- a/mitmproxy/optmanager.py
+++ b/mitmproxy/optmanager.py
@@ -3,6 +3,8 @@ import blinker
 import pprint
 import inspect
 import copy
+import functools
+import weakref
 
 from mitmproxy import exceptions
 from mitmproxy.utils import typecheck
@@ -63,6 +65,27 @@ class OptManager:
             self.__dict__["_opts"] = old
             self.changed.send(self, updated=updated)
 
+    def subscribe(self, func, opts):
+        """
+            Subscribe a callable to the .changed signal, but only for a
+            specified list of options. The callable should accept arguments
+            (options, updated), and may raise an OptionsError.
+        """
+        func = weakref.proxy(func)
+
+        @functools.wraps(func)
+        def _call(options, updated):
+            if updated.intersection(set(opts)):
+                try:
+                    func(options, updated)
+                except ReferenceError:
+                    self.changed.disconnect(_call)
+
+        # Our wrapper function goes out of scope immediately, so we have to set
+        # weakrefs to false. This means we need to keep our own weakref, and
+        # clean up the hook when it's gone.
+        self.changed.connect(_call, weak=False)
+
     def __eq__(self, other):
         return self._opts == other._opts
 
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 455824d3..27f5cb58 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -77,9 +77,6 @@ class ConsoleMaster(master.Master):
         self.options = self.options  # type: Options
         self.options.errored.connect(self.options_error)
 
-        self.palette = options.palette
-        self.palette_transparent = options.palette_transparent
-
         self.logbuffer = urwid.SimpleListWalker([])
 
         self.view_stack = []
@@ -253,10 +250,11 @@ class ConsoleMaster(master.Master):
         self.ui.start()
         os.unlink(name)
 
-    def set_palette(self, name):
-        self.palette = name
+    def set_palette(self, options, updated):
         self.ui.register_palette(
-            palettes.palettes[name].palette(self.palette_transparent)
+            palettes.palettes[options.palette].palette(
+                options.palette_transparent
+            )
         )
         self.ui.clear()
 
@@ -269,7 +267,11 @@ class ConsoleMaster(master.Master):
     def run(self):
         self.ui = urwid.raw_display.Screen()
         self.ui.set_terminal_properties(256)
-        self.set_palette(self.palette)
+        self.set_palette(self.options, None)
+        self.options.subscribe(
+            self.set_palette,
+            ["palette", "palette_transparent"]
+        )
         self.loop = urwid.MainLoop(
             urwid.SolidFill("x"),
             screen = self.ui,
diff --git a/mitmproxy/tools/console/palettepicker.py b/mitmproxy/tools/console/palettepicker.py
index a3eb9b90..0d943baf 100644
--- a/mitmproxy/tools/console/palettepicker.py
+++ b/mitmproxy/tools/console/palettepicker.py
@@ -3,7 +3,6 @@ import urwid
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import palettes
 from mitmproxy.tools.console import select
-from mitmproxy.tools.console import signals
 
 footer = [
     ('heading_key', "enter/space"), ":select",
@@ -43,8 +42,8 @@ class PalettePicker(urwid.WidgetWrap):
             return select.Option(
                 i,
                 None,
-                lambda: self.master.palette == name,
-                lambda: self.select(name)
+                lambda: self.master.options.palette == name,
+                lambda: setattr(self.master.options, "palette", name)
             )
 
         for i in high:
@@ -59,8 +58,8 @@ class PalettePicker(urwid.WidgetWrap):
                 select.Option(
                     "Transparent",
                     "T",
-                    lambda: master.palette_transparent,
-                    self.toggle_palette_transparent
+                    lambda: master.options.palette_transparent,
+                    master.options.toggler("palette_transparent")
                 )
             ]
         )
@@ -73,15 +72,7 @@ class PalettePicker(urwid.WidgetWrap):
             self.lb,
             header = title
         )
-        signals.update_settings.connect(self.sig_update_settings)
+        master.options.changed.connect(self.sig_options_changed)
 
-    def sig_update_settings(self, sender):
+    def sig_options_changed(self, options, updated):
         self.lb.walker._modified()
-
-    def select(self, name):
-        self.master.set_palette(name)
-
-    def toggle_palette_transparent(self):
-        self.master.palette_transparent = not self.master.palette_transparent
-        self.master.set_palette(self.master.palette)
-        signals.update_settings.send(self)
diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py
index 345512fd..c7808391 100644
--- a/test/mitmproxy/test_optmanager.py
+++ b/test/mitmproxy/test_optmanager.py
@@ -104,6 +104,29 @@ def test_toggler():
         o.toggler("nonexistent")
 
 
+class Rec():
+    def __init__(self):
+        self.called = None
+
+    def __call__(self, *args, **kwargs):
+        self.called = (args, kwargs)
+
+
+def test_subscribe():
+    o = TO()
+    r = Rec()
+    o.subscribe(r, ["two"])
+    o.one = "foo"
+    assert not r.called
+    o.two = "foo"
+    assert r.called
+
+    assert len(o.changed.receivers) == 1
+    del r
+    o.two = "bar"
+    assert len(o.changed.receivers) == 0
+
+
 def test_rollback():
     o = TO(one="two")