finalize script reloading :tada:

5 files changed, 95 insertions, 58 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 6f09a6b3..cef2013e 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -732,6 +732,8 @@ class ConsoleMaster(flow.FlowMaster):
             self.process_flow(f)
         return f
 
-    def script_change(self, script):
-        self.masterq.put(("script_change", script))
-        signals.status_message.send(message="<{}> reloaded.".format(script.args[0]))
+    def handle_script_change(self, script):
+        if super(ConsoleMaster, self).handle_script_change(script):
+            signals.status_message.send(message='"{}" reloaded.'.format(script.filename))
+        else:
+            signals.status_message.send(message='Error reloading "{}".'.format(script.filename))
\ No newline at end of file
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index f090d9c6..97d72992 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -640,7 +640,6 @@ class FlowMaster(controller.Master):
 
         self.stream = None
         self.apps = AppRegistry()
-        script.script_change.connect(self.script_change)
 
     def start_app(self, host, port):
         self.apps.add(
@@ -664,15 +663,18 @@ class FlowMaster(controller.Master):
             script_obj.unload()
         except script.ScriptException as e:
             self.add_event("Script error:\n" + str(e), "error")
+        script.reloader.unwatch(script_obj)
         self.scripts.remove(script_obj)
     
-    def load_script(self, command):
+    def load_script(self, command, use_reloader=True):
         """
             Loads a script. Returns an error description if something went
             wrong.
         """
         try:
             s = script.Script(command, script.ScriptContext(self))
+            if use_reloader:
+                script.reloader.watch(s, lambda: self.masterq.put(("script_change", s)))
         except script.ScriptException as v:
             return v.args[0]
         self.scripts.append(s)
@@ -1020,8 +1022,33 @@ class FlowMaster(controller.Master):
     def handle_accept_intercept(self, f):
         self.state.update_flow(f)
 
-    def handle_script_change(self, script):
-        script.load()
+    def handle_script_change(self, s):
+        """
+        Handle a script whose contents have been changed on the file system.
+
+        Args:
+            s (script.Script): the changed script
+
+        Returns:
+            True, if reloading was successful.
+            False, otherwise.
+        """
+        ok = True
+        # We deliberately do not want to fail here.
+        # In the worst case, we have an "empty" script object.
+        try:
+            s.unload()
+        except script.ScriptException as e:
+            ok = False
+            self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)))
+        try:
+            s.load()
+        except script.ScriptException as e:
+            ok = False
+            self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)))
+        else:
+            self.add_event('"{}" reloaded.'.format(s.filename))
+        return ok
 
     def shutdown(self):
         self.unload_scripts()
@@ -1039,11 +1066,6 @@ class FlowMaster(controller.Master):
         self.stream.fo.close()
         self.stream = None
 
-    def script_change(self, script):
-        self.masterq.put(("script_change", script))
-        self.add_event("<{}> reloaded.".format(script.args[0]))
-
-
 def read_flows_from_paths(paths):
     """
     Given a list of filepaths, read all flows and return a list of them.
diff --git a/libmproxy/script/__init__.py b/libmproxy/script/__init__.py
index 0f487795..8bcdc5a2 100644
--- a/libmproxy/script/__init__.py
+++ b/libmproxy/script/__init__.py
@@ -1,11 +1,13 @@
-from .script import Script, script_change
+from .script import Script
 from .script_context import ScriptContext
 from .concurrent import concurrent
 from ..exceptions import ScriptException
+from . import reloader
 
 __all__ = [
-    "Script", "script_change",
+    "Script",
     "ScriptContext",
     "concurrent",
-    "ScriptException"
+    "ScriptException",
+    "reloader"
 ]
\ No newline at end of file
diff --git a/libmproxy/script/reloader.py b/libmproxy/script/reloader.py
new file mode 100644
index 00000000..b867238f
--- /dev/null
+++ b/libmproxy/script/reloader.py
@@ -0,0 +1,37 @@
+import os
+from watchdog.events import PatternMatchingEventHandler
+from watchdog.observers import Observer
+
+_observers = {}
+
+
+def watch(script, callback):
+    script_dir = os.path.dirname(os.path.abspath(script.args[0]))
+    event_handler = _ScriptModificationHandler(callback)
+    observer = Observer()
+    observer.schedule(event_handler, script_dir)
+    observer.start()
+    _observers[script] = observer
+
+
+def unwatch(script):
+    observer = _observers.pop(script, None)
+    if observer:
+        observer.stop()
+
+
+class _ScriptModificationHandler(PatternMatchingEventHandler):
+    def __init__(self, callback):
+        # We could enumerate all relevant *.py files (as werkzeug does it),
+        # but our case looks like it isn't as simple as enumerating sys.modules.
+        # This should be good enough for now.
+        super(_ScriptModificationHandler, self).__init__(
+            ignore_directories=True,
+            patterns=["*.py"]
+        )
+        self.callback = callback
+
+    def on_modified(self, event):
+        self.callback()
+
+__all__ = ["watch", "unwatch"]
\ No newline at end of file
diff --git a/libmproxy/script/script.py b/libmproxy/script/script.py
index a58ba0af..498caf94 100644
--- a/libmproxy/script/script.py
+++ b/libmproxy/script/script.py
@@ -8,32 +8,27 @@ import os
 import shlex
 import traceback
 import sys
-import blinker
-
-from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent
-from watchdog.observers import Observer
-
 from ..exceptions import ScriptException
 
-script_change = blinker.Signal()
-
 
 class Script(object):
     """
-        Script object representing an inline script.
+    Script object representing an inline script.
     """
 
-    def __init__(self, command, context, use_reloader=True):
+    def __init__(self, command, context):
         self.command = command
         self.args = self.parse_command(command)
         self.ctx = context
         self.ns = None
         self.load()
-        if use_reloader:
-            self.start_observe()
 
-    @classmethod
-    def parse_command(cls, command):
+    @property
+    def filename(self):
+        return self.args[0]
+
+    @staticmethod
+    def parse_command(command):
         if not command or not command.strip():
             raise ScriptException("Empty script command.")
         if os.name == "nt":  # Windows: escape all backslashes in the path.
@@ -64,21 +59,22 @@ class Script(object):
         if self.ns is not None:
             self.unload()
         script_dir = os.path.dirname(os.path.abspath(self.args[0]))
-        ns = {'__file__': os.path.abspath(self.args[0])}
+        self.ns = {'__file__': os.path.abspath(self.args[0])}
         sys.path.append(script_dir)
         try:
-            execfile(self.args[0], ns, ns)
+            execfile(self.args[0], self.ns, self.ns)
         except Exception as e:
             # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/
             raise ScriptException(traceback.format_exc(e))
-        sys.path.pop()
-        self.ns = ns
+        finally:
+            sys.path.pop()
         return self.run("start", self.args)
 
     def unload(self):
-        ret = self.run("done")
-        self.ns = None
-        return ret
+        try:
+            return self.run("done")
+        finally:
+            self.ns = None
 
     def run(self, name, *args, **kwargs):
         """
@@ -98,26 +94,4 @@ class Script(object):
             except Exception as e:
                 raise ScriptException(traceback.format_exc(e))
         else:
-            return None
-
-    def start_observe(self):
-        script_dir = os.path.dirname(self.args[0])
-        event_handler = ScriptModified(self)
-        observer = Observer()
-        observer.schedule(event_handler, script_dir)
-        observer.start()
-
-    def stop_observe(self):
-        raise NotImplementedError() # FIXME
-
-
-class ScriptModified(PatternMatchingEventHandler):
-    def __init__(self, script):
-        # We could enumerate all relevant *.py files (as werkzeug does it),
-        # but our case looks like it isn't as simple as enumerating sys.modules.
-        # This should be good enough for now.
-        super(ScriptModified, self).__init__(ignore_directories=True, patterns=["*.py"])
-        self.script = script
-
-    def on_modified(self, event=FileModifiedEvent):
-        script_change.send(self.script)
+            return None
\ No newline at end of file