fix saving

dalejung · Jun 26, 2023 · 5f098ac · 5f098ac
1 parent 1252fa8
commit 5f098ac
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 95 deletions.
diff --git a/nbx_deux/bundle_manager/bundle.py b/nbx_deux/bundle_manager/bundle.py
@@ -20,10 +20,12 @@
 import nbformat
 from IPython.utils import tz
 
-from nbx_deux.models import BaseModel
+from nbx_deux.models import BaseModel, NotebookModel
 from nbx_deux.fileio import (
     _read_notebook,
     _save_notebook,
+    check_and_sign,
+    ospath_is_writable,
 )
 
 
@@ -169,6 +171,7 @@ def save(self, model):
             os.mkdir(bundle_path)
 
         self.save_bundle_file(model)
+        return model
 
     def save_bundle_file(self, model):
         content = model['content']
@@ -195,23 +198,19 @@ def get_model(self, root_dir, content=True, file_content=None):
             file_content = content
 
         os_path = self.bundle_file
-        info = os.stat(os_path)
-        last_modified = tz.utcfromtimestamp(info.st_mtime)
-        created = tz.utcfromtimestamp(info.st_ctime)
 
         bundle_file_content = None
         if content:
             bundle_file_content = self.get_bundle_file_content()
 
+        model = BaseModel.from_filepath_dict(os_path, root_dir)
+        assert model['name'] == self.name
+
         files = self.files_pack(file_content)
-        path = os.path.relpath(self.bundle_path, root_dir)
         model = self.bundle_model_class(
-            name=self.name,
-            path=path,
-            last_modified=last_modified,
-            created=created,
             bundle_files=files,
             content=bundle_file_content,
+            **model,
         )
         return model
 
@@ -261,12 +260,13 @@ def valid_path(cls, os_path):
             return False
         return cls.is_bundle(os_path)
 
-    def save_bundle_file(self, model):
+    def save_bundle_file(self, model: NotebookModel):
         nb = cast(nbformat.NotebookNode, nbformat.from_dict(model['content']))
+        check_and_sign(nb)
 
-        # TODO: I don't remember why I did this...
-        if 'name' in nb.metadata:
-            nb.metadata['name'] = u''
+        # # TODO: I don't remember why I did this...
+        # if 'name' in nb.metadata:
+        #     nb.metadata['name'] = u''
 
         _save_notebook(self.bundle_file, nb)
 
@@ -280,31 +280,12 @@ def get_bundle_file_content(self):
     from nbformat.v4 import new_notebook, writes
 
     with TempDir() as td:
-        subdir = td.joinpath('subdir')
-        nb_dir = subdir.joinpath('example.ipynb')
-        nb_dir.mkdir(parents=True)
-        file1 = nb_dir.joinpath('howdy.txt')
-        with file1.open('w') as f:
-            f.write('howdy')
-
-        nb_file = nb_dir.joinpath('example.ipynb')
+        nb_dir = td.joinpath('example.ipynb')
         nb = new_notebook()
-        nb['metadata']['howdy'] = 'hi'
-        with nb_file.open('w') as f:
-            f.write(writes(nb))
-
-        nb_bundle = NotebookBundlePath(nb_dir)
-        files = nb_bundle.files
-        assert 'howdy.txt' in files
-
-        content = nb_bundle.get_bundle_file_content()
-        assert content == nb
-
-        model = nb_bundle.get_model(td)
-        assert model.is_bundle is True
-        assert model.content == nb
-        assert model.bundle_files['howdy.txt'] == 'howdy'
-        assert model.name == 'example.ipynb'
-        assert model.path == 'subdir/example.ipynb'
+        bundle = NotebookBundlePath(nb_dir)
+        model = NotebookModel.from_nbnode(nb, name='hi.ipynb', path='hi.ipynb')
+        bundle.save(model)
+        assert nb_dir.is_dir()
 
-        items = bundle_list_dir(subdir)
+        new_model = bundle.get_model(td)
+        assert new_model['content'] == nb
diff --git a/nbx_deux/bundle_manager/bundle_nbmanager.py b/nbx_deux/bundle_manager/bundle_nbmanager.py
@@ -10,7 +10,7 @@
 from IPython.utils import tz
 from jupyter_server.services.contents.filemanager import FileContentsManager
 
-from nbx_deux.models import DirectoryModel
+from nbx_deux.models import DirectoryModel, NotebookModel
 
 from ..nbx_manager import NBXContentsManager, ApiPath
 from .bundle import NotebookBundlePath, BundlePath, bundle_get_path_item
@@ -30,9 +30,12 @@ def is_bundle(self, path: ApiPath | Path):
             os_path = self._get_os_path(path=path)
         return BundlePath.valid_path(os_path)
 
+    def is_notebook(self, path: ApiPath, type=None):
+        return (type is None and path.endswith(".ipynb"))
+
     def get_bundle(self, path: ApiPath, type=None):
         os_path = self._get_os_path(path=path)
-        if type == "notebook" or (type is None and path.endswith(".ipynb")):
+        if type == "notebook" or self.is_notebook(path, type):
             bundle = NotebookBundlePath(os_path)
         else:
             bundle = BundlePath(os_path)
@@ -80,9 +83,17 @@ def bundle_get(self, path, content=True, type=None, format=None):
         return model
 
     def save(self, model, path):
-        if self.is_bundle(path):
+        os_path = self._get_os_path(path=path)
+        is_notebook = self.is_notebook(path)
+        is_new = not os.path.exists(os_path)
+        is_new_notebook = is_new and is_notebook
+        # new files default to bundle
+        if self.is_bundle(path) or is_new_notebook:
             bundle = self.get_bundle(path)
-            return bundle.save(model)
+            bundle.save(model)
+            # refresh
+            model = self.get(path, content=False)
+            return model.asdict()
 
         return self.fm.save(model, path)
 
@@ -263,56 +274,10 @@ def delete_checkpoint(self, checkpoint_id, path=''):
     with TempDir() as td:
         subdir = td.joinpath('subdir')
 
-        regular_nb = td.joinpath("regular.ipynb")
-        nb = new_notebook()
-        with regular_nb.open('w') as f:
-            f.write(writes(nb))
-
-        regular_file = td.joinpath("sup.txt")
-        with regular_file.open('w') as f:
-            f.write("sups")
-
-        nb_dir = subdir.joinpath('example.ipynb')
-        nb_dir.mkdir(parents=True)
-        file1 = nb_dir.joinpath('howdy.txt')
-        with file1.open('w') as f:
-            f.write('howdy')
+        nbm = BundleContentsManager(root_dir=str(td))
 
-        nb_file = nb_dir.joinpath('example.ipynb')
         nb = new_notebook()
         nb['metadata']['howdy'] = 'hi'
-        with nb_file.open('w') as f:
-            f.write(writes(nb))
-
-        # try a regular bundle
-        bundle_dir = td.joinpath('example.txt')
-        bundle_dir.mkdir(parents=True)
-        bundle_file = bundle_dir.joinpath('example.txt')
-        with bundle_file.open('w') as f:
-            f.write("regular ole bundle")
-
-        nbm = BundleContentsManager(root_dir=str(td))
-        model = nbm.get("")
-        contents_dict = model.contents_dict()
-        assert contents_dict['subdir']['type'] == 'directory'
-        assert contents_dict['subdir']['content'] is None
-
-        assert contents_dict['regular.ipynb']['type'] == 'notebook'
-        assert contents_dict['regular.ipynb']['content'] is None
-
-        assert contents_dict['sup.txt']['type'] == 'file'
-        assert contents_dict['sup.txt']['content'] is None
-
-        notebook_model = nbm.get("subdir/example.ipynb")
-
-        assert notebook_model['type'] == 'notebook'
-        assert notebook_model['is_bundle'] is True
-        assert notebook_model['bundle_files'] == {'howdy.txt': 'howdy'}
-
-        subdir_model = nbm.get("subdir")
-        subdir_contents_dict = subdir_model.contents_dict()
-        assert subdir_contents_dict['subdir/example.ipynb']['type'] == 'notebook'
-        assert subdir_contents_dict['subdir/example.ipynb']['is_bundle'] is True
+        model = NotebookModel.from_nbnode(nb, path='hi.ipynb', name='hi.ipynb')
 
-        nb_model = nbm.get("subdir/example.ipynb", content=False)
-        assert subdir_contents_dict['subdir/example.ipynb'] == nb_model
+        new_model = nbm.save(model, 'dale.ipynb')
diff --git a/nbx_deux/bundle_manager/tests/test_bundle.py b/nbx_deux/bundle_manager/tests/test_bundle.py
@@ -1,3 +1,4 @@
+from nbx_deux.models import NotebookModel
 from nbx_deux.testing import TempDir
 from nbformat.v4 import new_notebook, writes
 
@@ -48,3 +49,16 @@ def test_notebook_bundle_file():
 
         new_model = nb_bundle.get_model(td, file_content=False)
         assert new_model.bundle_files['howdy.txt'] is None
+
+
+def test_notebook_bundle_save():
+    with TempDir() as td:
+        nb_dir = td.joinpath('example.ipynb')
+        nb = new_notebook()
+        bundle = NotebookBundlePath(nb_dir)
+        model = NotebookModel.from_nbnode(nb, name='hi.ipynb', path='hi.ipynb')
+        bundle.save(model)
+        assert nb_dir.is_dir()
+
+        new_model = bundle.get_model(td)
+        assert new_model['content'] == nb
diff --git a/nbx_deux/fileio.py b/nbx_deux/fileio.py
@@ -25,6 +25,9 @@
 from jupyter_server import _tz as tz
 
 
+CM_NOTARY = cast(sign.NotebookNotary, FileContentsManager().notary)
+
+
 def mark_trusted_cells(nb):
     """Mark cells as trusted if the notebook signature matches.
 
@@ -38,7 +41,7 @@ def mark_trusted_cells(nb):
         The notebook's path (for logging)
     """
 
-    notary = sign.NotebookNotary()
+    notary = CM_NOTARY
     trusted = notary.check_signature(nb)
     notary.mark_cells(nb, trusted)
 
@@ -231,3 +234,20 @@ def validate_notebook_model(model, validation_error=None):
             json.dumps(e.instance, indent=1, default=lambda obj: "<UNKNOWN>"),
         )
     return model
+
+
+def check_and_sign(nb):
+    """Check for trusted cells, and sign the notebook.
+
+    Called as a part of saving notebooks.
+
+    Parameters
+    ----------
+    nb : dict
+        The notebook dict
+    path : str
+        The notebook's path (for logging)
+    """
+    notary = CM_NOTARY
+    if notary.check_cells(nb):
+        notary.sign(nb)
diff --git a/nbx_deux/models.py b/nbx_deux/models.py
@@ -15,6 +15,7 @@
 from jupyter_core.paths import is_file_hidden
 from jupyter_server.services.contents.filemanager import FileContentsManager
 from jupyter_server.utils import ApiPath, to_os_path
+from nbformat import NotebookNode
 
 from nbx_deux.fileio import (
     FCM_HIDE_GLOBS,
@@ -197,6 +198,9 @@ def items(self):
     def keys(self):
         return self.asdict().keys()
 
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
 
 @dc.dataclass(kw_only=True)
 class FileModel(BaseModel):
@@ -232,6 +236,14 @@ class NotebookModel(BaseModel):
     type: str = dc.field(default='notebook', init=False)
     default_format: ClassVar = 'json'
 
+    @classmethod
+    def from_nbnode(cls, nb: NotebookNode, **kwargs):
+        now = datetime.now()
+
+        created = kwargs.pop('created', now)
+        last_modified = kwargs.pop('last_modified', now)
+        return cls(content=nb, created=created, last_modified=last_modified, **kwargs)
+
     @classmethod
     def from_filepath_dict(cls, os_path, root_dir=None, content=True, format=None):
         model = BaseModel.from_filepath_dict(os_path, root_dir=root_dir)
@@ -241,8 +253,6 @@ def from_filepath_dict(cls, os_path, root_dir=None, content=True, format=None):
             nb = _read_notebook(os_path)
             mark_trusted_cells(nb)
             model["content"] = nb
-            # Copying jupyter idiom of only setting format when content is requested
-            model["format"] = "json"
             validate_notebook_model(model, validation_error)
 
         return model