Clean up and clarify script API

nderkach · Jan 12, 2014 · e5776b8 · e5776b8
1 parent 58e1b3a
commit e5776b8
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 30 deletions.
diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
@@ -446,12 +446,14 @@ def run_script_once(self, path, f):
         if not path:
             return
         self.add_event("Running script on flow: %s"%path)
-        ret = self.get_script(shlex.split(path, posix=(os.name != "nt")))
-        if ret[0]:
+
+        try:
+            s = script.Script(shlex.split(path, posix=(os.name != "nt")), self)
+        except script.ScriptError, v:
             self.statusbar.message("Error loading script.")
-            self.add_event("Error loading script:\n%s"%ret[0])
+            self.add_event("Error loading script:\n%s"%v.args[0])
             return
-        s = ret[1]
+
         if f.request:
             self._run_script_method("request", s, f)
         if f.response:

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
@@ -1394,17 +1394,6 @@ def add_event(self, e, level="info"):
         """
         pass
 
-    def get_script(self, script_argv):
-        """
-            Returns an (error, script) tuple.
-        """
-        s = script.Script(script_argv, self)
-        try:
-            s.load()
-        except script.ScriptError, v:
-            return (v.args[0], None)
-        return (None, s)
-
     def unload_script(self, script):
         script.unload()
         self.scripts.remove(script)
@@ -1414,11 +1403,11 @@ def load_script(self, script_argv):
             Loads a script. Returns an error description if something went
             wrong.
         """
-        r = self.get_script(script_argv)
-        if r[0]:
-            return r[0]
-        else:
-            self.scripts.append(r[1])
+        try:
+            s = script.Script(script_argv, self)
+        except script.ScriptError, v:
+            return v.args[0]
+        self.scripts.append(s)
 
     def run_single_script_hook(self, script, name, *args, **kwargs):
         if script and not self.pause_scripts:

diff --git a/libmproxy/script.py b/libmproxy/script.py
@@ -49,6 +49,7 @@ def __init__(self, argv, master):
         self.argv = argv
         self.ctx = ScriptContext(master)
         self.ns = None
+        self.load()
 
     def load(self):
         """

diff --git a/test/test_script.py b/test/test_script.py
@@ -41,28 +41,24 @@ def test_err(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
 
-        s = script.Script(["nonexistent"], fm)
         tutils.raises(
             "no such file",
-            s.load
+            script.Script, ["nonexistent"], fm
         )
 
-        s = script.Script([tutils.test_data.path("scripts")], fm)
         tutils.raises(
             "not a file",
-            s.load
+            script.Script, [tutils.test_data.path("scripts")], fm
         )
 
-        s = script.Script([tutils.test_data.path("scripts/syntaxerr.py")], fm)
         tutils.raises(
             script.ScriptError,
-            s.load
+            script.Script, [tutils.test_data.path("scripts/syntaxerr.py")], fm
         )
 
-        s = script.Script([tutils.test_data.path("scripts/loaderr.py")], fm)
         tutils.raises(
             script.ScriptError,
-            s.load
+            script.Script, [tutils.test_data.path("scripts/loaderr.py")], fm
         )
 
     def test_concurrent(self):
@@ -106,8 +102,7 @@ def test_concurrent2(self):
     def test_concurrent_err(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
-        s = script.Script([tutils.test_data.path("scripts/concurrent_decorator_err.py")], fm)
         tutils.raises(
             "decorator not supported for this method",
-            s.load
+            script.Script, [tutils.test_data.path("scripts/concurrent_decorator_err.py")], fm
         )