Repository: bokeh/bokeh

Issue:
Fix `curdoc()` inside of callbacks
When running a script, `ScriptHandler` sets curdoc here: 
https://github.com/bokeh/bokeh/blob/129594fda0fe5aa0645d38c6642f2d580ee55407/bokeh/application/spellings/script.py#L79

It does not unset it after running the script, which has multiple problems:
- curdoc() may appear to work in `on_change` callbacks as long as there's one session from one app, but will break in production
- security leakage where sessions see the wrong curdoc()

Generally what an `on_change` callback should do is look at `model.document` for some model, to get the right document.

Minimum fix: set curdoc() to None after we run the `ScriptHandler`
Fancier fix: add a feature to Document where it sets itself as `curdoc()` temporarily whenever it runs callbacks, then sets back the previous `curdoc()`



Discussion:

Resolving patch:
diff --git a/bokeh/application/spellings/script.py b/bokeh/application/spellings/script.py
--- a/bokeh/application/spellings/script.py
+++ b/bokeh/application/spellings/script.py
@@ -78,6 +78,7 @@ def modify_document(self, doc):
         self._module_name = 'bk_script_' + str(uuid.uuid4()).replace('-', '')
         self._module = ModuleType(self._module_name)
         self._module.__dict__['__file__'] = abspath(self._path)
+        old_doc = curdoc()
         set_curdoc(doc)
         old_io = self._monkeypatch_io()
         try:
@@ -97,3 +98,4 @@ def modify_document(self, doc):
             self._error = "%s\nFile \"%s\", line %d, in %s:\n%s" % (str(e), os.path.basename(filename), line_number, func, txt)
         finally:
             self._unmonkeypatch_io(old_io)
+            set_curdoc(old_doc)
diff --git a/bokeh/document.py b/bokeh/document.py
--- a/bokeh/document.py
+++ b/bokeh/document.py
@@ -388,9 +388,28 @@ def remove_on_change(self, *callbacks):
         for callback in callbacks:
             self._callbacks.remove(callback)
 
+    def _with_self_as_curdoc(self, f):
+        from bokeh.io import set_curdoc, curdoc
+        old_doc = curdoc()
+        try:
+            set_curdoc(self)
+            f()
+        finally:
+            set_curdoc(old_doc)
+
+    def _wrap_with_self_as_curdoc(self, f):
+        doc = self
+        def wrapper(*args, **kwargs):
+            def invoke():
+                f(*args, **kwargs)
+            doc._with_self_as_curdoc(invoke)
+        return wrapper
+
     def _trigger_on_change(self, event):
-        for cb in self._callbacks:
-            cb(event)
+        def invoke_callbacks():
+            for cb in self._callbacks:
+                cb(event)
+        self._with_self_as_curdoc(invoke_callbacks)
 
     def _notify_change(self, model, attr, old, new):
         ''' Called by PlotObject when it changes
@@ -657,7 +676,7 @@ def add_periodic_callback(self, callback, period, id=None):
 
         '''
         # create the new callback object
-        cb = PeriodicCallback(self, callback, period, id)
+        cb = PeriodicCallback(self, self._wrap_with_self_as_curdoc(callback), period, id)
         self._session_callbacks[callback] = cb
         # emit event so the session is notified of the new callback
         self._trigger_on_change(SessionCallbackAdded(self, cb))
@@ -678,7 +697,7 @@ def add_timeout_callback(self, callback, timeout, id=None):
 
         '''
         # create the new callback object
-        cb = TimeoutCallback(self, callback, timeout, id)
+        cb = TimeoutCallback(self, self._wrap_with_self_as_curdoc(callback), timeout, id)
         self._session_callbacks[callback] = cb
         # emit event so the session is notified of the new callback
         self._trigger_on_change(SessionCallbackAdded(self, cb))
diff --git a/bokeh/util/callback_manager.py b/bokeh/util/callback_manager.py
--- a/bokeh/util/callback_manager.py
+++ b/bokeh/util/callback_manager.py
@@ -86,9 +86,13 @@ def trigger(self, attr, old, new):
             None
 
         '''
+        def invoke():
+            callbacks = self._callbacks.get(attr)
+            if callbacks:
+                for callback in callbacks:
+                    callback(attr, old, new)
         if hasattr(self, '_document') and self._document is not None:
             self._document._notify_change(self, attr, old, new)
-        callbacks = self._callbacks.get(attr)
-        if callbacks:
-            for callback in callbacks:
-                callback(attr, old, new)
+            self._document._with_self_as_curdoc(invoke)
+        else:
+            invoke()


Test patch:
diff --git a/bokeh/tests/test_client_server.py b/bokeh/tests/test_client_server.py
--- a/bokeh/tests/test_client_server.py
+++ b/bokeh/tests/test_client_server.py
@@ -316,7 +316,6 @@ def cb(): pass
             for ss in [server_session, client_session, server_session2]:
                 iocb = ss._callbacks[callback.id]
                 assert isinstance(iocb, PeriodicCallback)
-                assert iocb.callback == cb
                 assert iocb.callback_time == 1
                 assert iocb.io_loop == server.io_loop
                 assert iocb.is_running()
diff --git a/bokeh/tests/test_document.py b/bokeh/tests/test_document.py
--- a/bokeh/tests/test_document.py
+++ b/bokeh/tests/test_document.py
@@ -3,6 +3,7 @@
 import unittest
 
 import bokeh.document as document
+from bokeh.io import curdoc
 from bokeh.plot_object import PlotObject
 from bokeh.properties import Int, Instance, String
 
@@ -254,8 +255,11 @@ def test_change_notification(self):
         d.add_root(m)
         assert len(d.roots) == 1
         assert m.bar == 1
+        assert curdoc() is not d
         events = []
+        curdoc_from_listener = []
         def listener(event):
+            curdoc_from_listener.append(curdoc())
             events.append(event)
         d.on_change(listener)
         m.bar = 42
@@ -267,6 +271,8 @@ def listener(event):
         assert event.attr == 'bar'
         assert event.old == 1
         assert event.new == 42
+        assert len(curdoc_from_listener) == 1
+        assert curdoc_from_listener[0] is d
 
     def test_change_notification_removal(self):
         d = document.Document()
@@ -356,7 +362,6 @@ def cb(): pass
         assert callback == d.session_callbacks[0] == events[0].callback
         assert callback.id == 'abc'
         assert callback.period == 1
-        assert callback.callback == cb
 
         callback = d.remove_periodic_callback(cb)
         assert len(d.session_callbacks) == 0
@@ -383,7 +388,6 @@ def cb(): pass
         assert callback == d.session_callbacks[0] == events[0].callback
         assert callback.id == 'abc'
         assert callback.timeout == 1
-        assert callback.callback == cb
 
         callback = d.remove_timeout_callback(cb)
         assert len(d.session_callbacks) == 0
@@ -391,6 +395,41 @@ def cb(): pass
         assert isinstance(events[0], document.SessionCallbackAdded)
         assert isinstance(events[1], document.SessionCallbackRemoved)
 
+    def test_periodic_callback_gets_curdoc(self):
+        d = document.Document()
+        assert curdoc() is not d
+        curdoc_from_cb = []
+        def cb():
+            curdoc_from_cb.append(curdoc())
+        callback = d.add_periodic_callback(cb, 1)
+        callback.callback()
+        assert len(curdoc_from_cb) == 1
+        assert curdoc_from_cb[0] is d
+
+    def test_timeout_callback_gets_curdoc(self):
+        d = document.Document()
+        assert curdoc() is not d
+        curdoc_from_cb = []
+        def cb():
+            curdoc_from_cb.append(curdoc())
+        callback = d.add_timeout_callback(cb, 1)
+        callback.callback()
+        assert len(curdoc_from_cb) == 1
+        assert curdoc_from_cb[0] is d
+
+    def test_model_callback_gets_curdoc(self):
+        d = document.Document()
+        m = AnotherModelInTestDocument(bar=42)
+        d.add_root(m)
+        assert curdoc() is not d
+        curdoc_from_cb = []
+        def cb(attr, old, new):
+            curdoc_from_cb.append(curdoc())
+        m.on_change('bar', cb)
+        m.bar = 43
+        assert len(curdoc_from_cb) == 1
+        assert curdoc_from_cb[0] is d
+
     def test_clear(self):
         d = document.Document()
         assert not d.roots