--- a/pydershell/pydershell.py	Mon Sep 07 04:38:44 2009 -0700
+++ b/pydershell/pydershell.py	Mon Sep 07 12:57:20 2009 -0700
@@ -71,32 +71,25 @@
     Securely wraps a JS object to behave like any normal Python object.
     """
 
-    __slots__ = ['_jsobject', '_cx', '_this']
+    __slots__ = ['_jsobject', '_sandbox', '_this']
 
-    def __init__(self, cx, jsobject, this):
+    def __init__(self, sandbox, jsobject, this):
         if not isinstance(jsobject, pydermonkey.Object):
             raise TypeError("Cannot wrap '%s' object" %
                             type(jsobject).__name__)
+        object.__setattr__(self, '_sandbox', sandbox)
         object.__setattr__(self, '_jsobject', jsobject)
-        object.__setattr__(self, '_cx', cx)
         object.__setattr__(self, '_this', this)
 
-    @staticmethod
-    def wrap(cx, jsvalue, this):
-        if isinstance(jsvalue, pydermonkey.Function):
-            return SafeJsFunctionWrapper(cx, jsvalue, this)
-        elif isinstance(jsvalue, pydermonkey.Object):
-            return SafeJsObjectWrapper(cx, jsvalue, this)
-        else:
-            # It's a primitive value.
-            return jsvalue
+    @property
+    def wrapped_jsobject(self):
+        return self._jsobject
 
     def _wrap_to_python(self, jsvalue):
-        return self.wrap(self._cx, jsvalue, self._jsobject)
+        return self._sandbox.wrap_jsobject(jsvalue, self._jsobject)
 
     def _wrap_to_js(self, value):
-        # TODO: Add support for wrapping non-primitive python objects.
-        return value
+        return self._sandbox.wrap_pyobject(value)
 
     def __eq__(self, other):
         if isinstance(other, SafeJsObjectWrapper):
@@ -114,7 +107,7 @@
         self.__setattr__(item, value)
 
     def __setattr__(self, name, value):
-        cx = self._cx
+        cx = self._sandbox.cx
         jsobject = self._jsobject
 
         cx.define_property(jsobject, name,
@@ -124,19 +117,19 @@
         return self.__getattr__(item)
 
     def __getattr__(self, name):
-        cx = self._cx
+        cx = self._sandbox.cx
         jsobject = self._jsobject
 
         return self._wrap_to_python(cx.get_property(jsobject, name))
 
     def __contains__(self, item):
-        cx = self._cx
+        cx = self._sandbox.cx
         jsobject = self._jsobject
 
         return cx.has_property(jsobject, item)
 
     def __iter__(self):
-        cx = self._cx
+        cx = self._sandbox.cx
         jsobject = self._jsobject
 
         properties = cx.enumerate(jsobject)
@@ -148,48 +141,23 @@
     Securely wraps a JS function to behave like any normal Python object.
     """
 
-    def __init__(self, cx, jsfunction, this):
+    def __init__(self, sandbox, jsfunction, this):
         if not isinstance(jsfunction, pydermonkey.Function):
             raise TypeError("Cannot wrap '%s' object" %
                             type(jsobject).__name__)
-        SafeJsObjectWrapper.__init__(self, cx, jsfunction, this)
+        SafeJsObjectWrapper.__init__(self, sandbox, jsfunction, this)
 
     def __call__(self, *args):
-        cx = self._cx
+        cx = self._sandbox.cx
         jsobject = self._jsobject
         this = self._this
 
-        obj = cx.call_function(this, jsobject,
-                               self._wrap_to_js(args))
-        return self._wrap_to_python(obj)
-
-def safejsfunc(cx, on_obj, name=None):
-    """
-    Exposes the decorated Python function on the given JS object.
-
-    Any unexpected exceptions raised by the function will be
-    re-raised as InternalError exceptions.
-    """
+        arglist = []
+        for arg in args:
+            arglist.append(self._wrap_to_js(arg))
 
-    def make_wrapper(func):
-        if name is None:
-            func_name = func.__name__
-        else:
-            func_name = name
-        def wrapper(func_cx, this, args):
-            try:
-                return func(func_cx, this, args)
-            except pydermonkey.error:
-                raise
-            except Exception:
-                raise InternalError()
-        cx.define_property(
-            on_obj,
-            func_name, 
-            cx.new_function(wrapper, func_name)
-            )
-        return func
-    return make_wrapper
+        obj = cx.call_function(this, jsobject, tuple(arglist))
+        return self._wrap_to_python(obj)
 
 def format_stack(js_stack):
     """
@@ -225,30 +193,17 @@
     lines.insert(0, "Traceback (most recent call last):")
     return '\n'.join(lines)
 
-class JSRuntime(object):
+class JSSandbox(object):
     """
-    A JS runtime capable of loading and executing scripts.
+    A JS runtime and associated functionality capable of securely
+    loading and executing scripts.
     """
 
     def __init__(self, watchdog=watchdog):
         rt = pydermonkey.Runtime()
         cx = rt.new_context()
-        globalobj = cx.new_object()
-        cx.init_standard_classes(globalobj)
-
-        @safejsfunc(cx, globalobj)
-        def bar(cx, this, args):
-            obj = SafeJsObjectWrapper.wrap(cx, args[0], this)
-            print obj.bar()
-
-        @safejsfunc(cx, globalobj)
-        def foo(cx, this, args):
-            return cx.call_function(this, args[0], ())
-
-        @safejsfunc(cx, globalobj, 'print')
-        def jsprint(cx, this, args):
-            if len(args) > 0:
-                print args[0]
+        root = cx.new_object()
+        cx.init_standard_classes(root)
 
         cx.set_operation_callback(self._opcb)
         cx.set_throw_hook(self._throwhook)
@@ -256,10 +211,21 @@
 
         self.rt = rt
         self.cx = cx
-        self.globalobj = globalobj
         self.curr_exc = None
         self.py_stack = None
         self.js_stack = None
+        self.__py_to_js = {}
+        self.root = self.wrap_jsobject(root, root)
+
+    def finish(self):
+        for jsobj in self.__py_to_js.values():
+            self.cx.clear_object_private(jsobj)
+        del self.__py_to_js
+        del self.curr_exc
+        del self.py_stack
+        del self.js_stack
+        del self.cx
+        del self.rt
 
     def _opcb(self, cx):
         # Don't do anything; if a keyboard interrupt was triggered,
@@ -273,11 +239,66 @@
             self.py_stack = traceback.extract_stack()
             self.js_stack = cx.get_stack()
 
+    def __wrap_pycallable(self, func):
+        if func in self.__py_to_js:
+            return self.__py_to_js[func]
+
+        if hasattr(func, '__name__'):
+            name = func.__name__
+        else:
+            name = ""
+
+        def wrapper(func_cx, this, args):
+            try:
+                arglist = []
+                for arg in args:
+                    arglist.append(self.wrap_jsobject(arg))
+                # TODO: Fill in extra required params with pymonkey.undefined?
+                # or automatically throw an exception to calling js code?
+                return func(*arglist)
+            except pydermonkey.error:
+                raise
+            except Exception:
+                raise InternalError()
+        wrapper.wrapped_pyobject = func
+        wrapper.__name__ = name
+
+        jsfunc = self.cx.new_function(wrapper, name)
+        self.__py_to_js[func] = jsfunc
+
+        return jsfunc
+
+    def wrap_pyobject(self, value):
+        if isinstance(value, SafeJsObjectWrapper):
+            # It's already wrapped, just unwrap it.
+            return value.wrapped_jsobject
+        elif callable(value):
+            return self.__wrap_pycallable(value)
+        elif isinstance(value, object):
+            raise NotImplementedError("TODO")
+        else:
+            # Here's hoping it's a primitive value.
+            return value
+
+    def wrap_jsobject(self, jsvalue, this=None):
+        if this is None:
+            this = self.root.wrapped_jsobject
+        if isinstance(jsvalue, pydermonkey.Function):
+            if jsvalue.is_python:
+                # It's a Python function, just unwrap it.
+                return self.cx.get_object_private(jsvalue).wrapped_pyobject
+            return SafeJsFunctionWrapper(self, jsvalue, this)
+        elif isinstance(jsvalue, pydermonkey.Object):
+            return SafeJsObjectWrapper(self, jsvalue, this)
+        else:
+            # It's a primitive value.
+            return jsvalue
+
     def run_script(self, filename):
         contents = open(filename).read()
         cx = self.cx
         try:
-            cx.evaluate_script(self.globalobj, contents,
+            cx.evaluate_script(self.root.wrapped_jsobject, contents,
                                filename, 1)
         except pydermonkey.error, e:
             print format_stack(self.js_stack)
@@ -288,9 +309,24 @@
             print e.exc_info[1]
 
 if __name__ == '__main__':
-    runtime = JSRuntime()
-    runtime.run_script('test.js')
-    del runtime
+    sandbox = JSSandbox()
+
+    def bar(obj):
+        print obj.bar()
+    sandbox.root.bar = bar
+
+    def foo(callback):
+        return callback()
+    sandbox.root.foo = foo
+
+    def jsprint(string):
+        print string
+    sandbox.root['print'] = jsprint
+
+    sandbox.run_script('test.js')
+
+    sandbox.finish()
+    del sandbox
 
     import gc
     gc.collect()