1 files changed, 36 insertions, 9 deletions
diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index ab068e47..ed5f2ecd 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -16,6 +16,19 @@ import watchdog.events
 from watchdog.observers import polling
 
 
+class NS:
+    def __init__(self, ns):
+        self.__dict__["ns"] = ns
+
+    def __getattr__(self, key):
+        if key not in self.ns:
+            raise AttributeError("No such element: %s", key)
+        return self.ns[key]
+
+    def __setattr__(self, key, value):
+        self.__dict__["ns"][key] = value
+
+
 def parse_command(command):
     """
         Returns a (path, args) tuple.
@@ -74,18 +87,27 @@ def load_script(path, args):
     ns = {'__file__': os.path.abspath(path)}
     with scriptenv(path, args):
         exec(code, ns, ns)
-    return ns
+    return NS(ns)
 
 
 class ReloadHandler(watchdog.events.FileSystemEventHandler):
     def __init__(self, callback):
         self.callback = callback
 
+    def filter(self, event):
+        if event.is_directory:
+            return False
+        if os.path.basename(event.src_path).startswith("."):
+            return False
+        return True
+
     def on_modified(self, event):
-        self.callback()
+        if self.filter(event):
+            self.callback()
 
     def on_created(self, event):
-        self.callback()
+        if self.filter(event):
+            self.callback()
 
 
 class Script:
@@ -118,27 +140,32 @@ class Script:
         # It's possible for ns to be un-initialised if we failed during
         # configure
         if self.ns is not None and not self.dead:
-            func = self.ns.get(name)
+            func = getattr(self.ns, name, None)
             if func:
                 with scriptenv(self.path, self.args):
-                    func(*args, **kwargs)
+                    return func(*args, **kwargs)
 
     def reload(self):
         self.should_reload.set()
 
+    def load_script(self):
+        self.ns = load_script(self.path, self.args)
+        ret = self.run("start")
+        if ret:
+            self.ns = ret
+            self.run("start")
+
     def tick(self):
         if self.should_reload.is_set():
             self.should_reload.clear()
             ctx.log.info("Reloading script: %s" % self.name)
-            self.ns = load_script(self.path, self.args)
-            self.start()
+            self.load_script()
             self.configure(self.last_options)
         else:
             self.run("tick")
 
     def start(self):
-        self.ns = load_script(self.path, self.args)
-        self.run("start")
+        self.load_script()
 
     def configure(self, options):
         self.last_options = options