Register pre/post save hooks, call them sequentially (#696)

Author: davidbrochart

jupyter_server/services/contents/filemanager.py

@@ -11,13 +11,11 @@
 
 import nbformat
 from anyio.to_thread import run_sync
-from ipython_genutils.importstring import import_item
 from jupyter_core.paths import exists
 from jupyter_core.paths import is_file_hidden
 from jupyter_core.paths import is_hidden
 from send2trash import send2trash
 from tornado import web
-from traitlets import Any
 from traitlets import Bool
 from traitlets import default
 from traitlets import TraitError
@@ -54,48 +52,6 @@ def _default_root_dir(self):
         except AttributeError:
             return os.getcwd()
 
-    post_save_hook = Any(
-        None,
-        config=True,
-        allow_none=True,
-        help="""Python callable or importstring thereof
-
-        to be called on the path of a file just saved.
-
-        This can be used to process the file on disk,
-        such as converting the notebook to a script or HTML via nbconvert.
-
-        It will be called as (all arguments passed by keyword)::
-
-            hook(os_path=os_path, model=model, contents_manager=instance)
-
-        - path: the filesystem path to the file just written
-        - model: the model representing the file
-        - contents_manager: this ContentsManager instance
-        """,
-    )
-
-    @validate("post_save_hook")
-    def _validate_post_save_hook(self, proposal):
-        value = proposal["value"]
-        if isinstance(value, str):
-            value = import_item(value)
-        if not callable(value):
-            raise TraitError("post_save_hook must be callable")
-        return value
-
-    def run_post_save_hook(self, model, os_path):
-        """Run the post-save hook if defined, and log errors"""
-        if self.post_save_hook:
-            try:
-                self.log.debug("Running post-save hook on %s", os_path)
-                self.post_save_hook(os_path=os_path, model=model, contents_manager=self)
-            except Exception as e:
-                self.log.error("Post-save hook failed o-n %s", os_path, exc_info=True)
-                raise web.HTTPError(
-                    500, "Unexpected error while running post hook save: %s" % e
-                ) from e
-
     @validate("root_dir")
     def _validate_root_dir(self, proposal):
         """Do a bit of validation of the root_dir."""
@@ -451,7 +407,7 @@ def save(self, model, path=""):
         """Save the file model and return the model with no content."""
         path = path.strip("/")
 
-        self.run_pre_save_hook(model=model, path=path)
+        self.run_pre_save_hooks(model=model, path=path)
 
         if "type" not in model:
             raise web.HTTPError(400, "No file type provided")
@@ -491,7 +447,7 @@ def save(self, model, path=""):
         if validation_message:
             model["message"] = validation_message
 
-        self.run_post_save_hook(model=model, os_path=os_path)
+        self.run_post_save_hooks(model=model, os_path=os_path)
 
         return model
 
@@ -815,7 +771,7 @@ async def save(self, model, path=""):
         if validation_message:
             model["message"] = validation_message
 
-        self.run_post_save_hook(model=model, os_path=os_path)
+        self.run_post_save_hooks(model=model, os_path=os_path)
 
         return model
 

jupyter_server/services/contents/largefilemanager.py

@@ -56,7 +56,7 @@ def save(self, model, path=""):
 
             # Last chunk
             if chunk == -1:
-                self.run_post_save_hook(model=model, os_path=os_path)
+                self.run_post_save_hooks(model=model, os_path=os_path)
             return model
         else:
             return super(LargeFileManager, self).save(model, path)
@@ -131,7 +131,7 @@ async def save(self, model, path=""):
 
             # Last chunk
             if chunk == -1:
-                self.run_post_save_hook(model=model, os_path=os_path)
+                self.run_post_save_hooks(model=model, os_path=os_path)
             return model
         else:
             return await super(AsyncLargeFileManager, self).save(model, path)

jupyter_server/services/contents/manager.py

@@ -4,6 +4,7 @@
 import itertools
 import json
 import re
+import warnings
 from fnmatch import fnmatch
 
 from ipython_genutils.importstring import import_item
@@ -126,10 +127,55 @@ def _validate_pre_save_hook(self, proposal):
             value = import_item(self.pre_save_hook)
         if not callable(value):
             raise TraitError("pre_save_hook must be callable")
+        if self.pre_save_hook is not None:
+            warnings.warn(
+                f"Overriding existing pre_save_hook ({self.pre_save_hook.__name__}) with a new one ({value.__name__}).",
+                stacklevel=2,
+            )
+        return value
+
+    post_save_hook = Any(
+        None,
+        config=True,
+        allow_none=True,
+        help="""Python callable or importstring thereof
+
+        to be called on the path of a file just saved.
+
+        This can be used to process the file on disk,
+        such as converting the notebook to a script or HTML via nbconvert.
+
+        It will be called as (all arguments passed by keyword)::
+
+            hook(os_path=os_path, model=model, contents_manager=instance)
+
+        - path: the filesystem path to the file just written
+        - model: the model representing the file
+        - contents_manager: this ContentsManager instance
+        """,
+    )
+
+    @validate("post_save_hook")
+    def _validate_post_save_hook(self, proposal):
+        value = proposal["value"]
+        if isinstance(value, str):
+            value = import_item(value)
+        if not callable(value):
+            raise TraitError("post_save_hook must be callable")
+        if self.post_save_hook is not None:
+            warnings.warn(
+                f"Overriding existing post_save_hook ({self.post_save_hook.__name__}) with a new one ({value.__name__}).",
+                stacklevel=2,
+            )
         return value
 
     def run_pre_save_hook(self, model, path, **kwargs):
         """Run the pre-save hook if defined, and log errors"""
+        warnings.warn(
+            "run_pre_save_hook is deprecated, use run_pre_save_hooks instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         if self.pre_save_hook:
             try:
                 self.log.debug("Running pre-save hook on %s", path)
@@ -143,6 +189,77 @@ def run_pre_save_hook(self, model, path, **kwargs):
                 # which could cause frustrating data loss
                 self.log.error("Pre-save hook failed on %s", path, exc_info=True)
 
+    def run_post_save_hook(self, model, os_path):
+        """Run the post-save hook if defined, and log errors"""
+        warnings.warn(
+            "run_post_save_hook is deprecated, use run_post_save_hooks instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        if self.post_save_hook:
+            try:
+                self.log.debug("Running post-save hook on %s", os_path)
+                self.post_save_hook(os_path=os_path, model=model, contents_manager=self)
+            except Exception as e:
+                self.log.error("Post-save hook failed o-n %s", os_path, exc_info=True)
+                raise HTTPError(500, "Unexpected error while running post hook save: %s" % e) from e
+
+    _pre_save_hooks = List()
+    _post_save_hooks = List()
+
+    def register_pre_save_hook(self, hook):
+        if isinstance(hook, str):
+            hook = import_item(hook)
+        if not callable(hook):
+            raise RuntimeError("hook must be callable")
+        self._pre_save_hooks.append(hook)
+
+    def register_post_save_hook(self, hook):
+        if isinstance(hook, str):
+            hook = import_item(hook)
+        if not callable(hook):
+            raise RuntimeError("hook must be callable")
+        self._post_save_hooks.append(hook)
+
+    def run_pre_save_hooks(self, model, path, **kwargs):
+        """Run the pre-save hooks if any, and log errors"""
+        pre_save_hooks = [self.pre_save_hook] if self.pre_save_hook is not None else []
+        pre_save_hooks += self._pre_save_hooks
+        for pre_save_hook in pre_save_hooks:
+            try:
+                self.log.debug("Running pre-save hook on %s", path)
+                pre_save_hook(model=model, path=path, contents_manager=self, **kwargs)
+            except HTTPError:
+                # allow custom HTTPErrors to raise,
+                # rejecting the save with a message.
+                raise
+            except Exception:
+                # unhandled errors don't prevent saving,
+                # which could cause frustrating data loss
+                self.log.error(
+                    "Pre-save hook %s failed on %s",
+                    pre_save_hook.__name__,
+                    path,
+                    exc_info=True,
+                )
+
+    def run_post_save_hooks(self, model, os_path):
+        """Run the post-save hooks if any, and log errors"""
+        post_save_hooks = [self.post_save_hook] if self.post_save_hook is not None else []
+        post_save_hooks += self._post_save_hooks
+        for post_save_hook in post_save_hooks:
+            try:
+                self.log.debug("Running post-save hook on %s", os_path)
+                post_save_hook(os_path=os_path, model=model, contents_manager=self)
+            except Exception as e:
+                self.log.error(
+                    "Post-save %s hook failed on %s",
+                    post_save_hook.__name__,
+                    os_path,
+                    exc_info=True,
+                )
+                raise HTTPError(500, "Unexpected error while running post hook save: %s" % e) from e
+
     checkpoints_class = Type(Checkpoints, config=True)
     checkpoints = Instance(Checkpoints, config=True)
     checkpoints_kwargs = Dict(config=True)

tests/test_files.py

@@ -1,3 +1,4 @@
+import json
 import os
 from pathlib import Path
 
@@ -58,7 +59,8 @@ async def test_hidden_files(jp_fetch, jp_serverapp, jp_root_dir, maybe_hidden):
 
 
 async def test_contents_manager(jp_fetch, jp_serverapp, jp_root_dir):
-    """make sure ContentsManager returns right files (ipynb, bin, txt)."""
+    """make sure ContentsManager returns right files (ipynb, bin, txt).
+    Also test save file hooks."""
     nb = new_notebook(
         cells=[
             new_markdown_cell("Created by test ³"),
@@ -90,6 +92,48 @@ async def test_contents_manager(jp_fetch, jp_serverapp, jp_root_dir):
     assert r.body.decode() == "foobar"
 
 
+async def test_save_hooks(jp_fetch, jp_serverapp):
+    # define a first pre-save hook that will change the content of the file before saving
+    def pre_save_hook1(model, **kwargs):
+        model["content"] += " was modified"
+
+    # define a second pre-save hook that will change the content of the file before saving
+    def pre_save_hook2(model, **kwargs):
+        model["content"] += " twice!"
+
+    # define a first post-save hook that will change the 'last_modified' date
+    def post_save_hook1(model, **kwargs):
+        model["last_modified"] = "yesterday"
+
+    # define a second post-save hook that will change the 'last_modified' date
+    def post_save_hook2(model, **kwargs):
+        model["last_modified"] += " or tomorrow!"
+
+    # register the pre-save hooks
+    jp_serverapp.contents_manager.register_pre_save_hook(pre_save_hook1)
+    jp_serverapp.contents_manager.register_pre_save_hook(pre_save_hook2)
+
+    # register the post-save hooks
+    jp_serverapp.contents_manager.register_post_save_hook(post_save_hook1)
+    jp_serverapp.contents_manager.register_post_save_hook(post_save_hook2)
+
+    # send a request to save a file, with an original content
+    # the 'last_modified' returned model field should have been modified by post_save_hook1 then post_save_hook2
+    r = await jp_fetch(
+        "api/contents/test.txt",
+        method="PUT",
+        body=json.dumps(
+            {"format": "text", "path": "test.txt", "type": "file", "content": "original content"}
+        ),
+    )
+    assert json.loads(r.body.decode())["last_modified"] == "yesterday or tomorrow!"
+
+    # read the file back
+    # the original content should have been modified by pre_save_hook1 then pre_save_hook2
+    r = await jp_fetch("files/test.txt", method="GET")
+    assert r.body.decode() == "original content was modified twice!"
+
+
 async def test_download(jp_fetch, jp_serverapp, jp_root_dir):
     text = "hello"
     jp_root_dir.joinpath("test.txt").write_text(text)