lamindb 1.0.4__py3-none-any.whl → 1.1.0__py3-none-any.whl

lamindb/_finish.py

@@ -1,22 +1,84 @@
 from __future__ import annotations
 
+import builtins
 import re
 from datetime import datetime, timezone
+from time import sleep
 from typing import TYPE_CHECKING
 
 import lamindb_setup as ln_setup
 from lamin_utils import logger
+from lamin_utils._logger import LEVEL_TO_COLORS, LEVEL_TO_ICONS, RESET_COLOR
 from lamindb_setup.core.hashing import hash_file
 
-from lamindb.core.exceptions import NotebookNotSaved
 from lamindb.models import Artifact, Run, Transform
 
+is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
+
 if TYPE_CHECKING:
     from pathlib import Path
 
 
 def get_save_notebook_message() -> str:
-    return f"Please save the notebook in your editor (shortcut `{get_shortcut()}`) within 2 sec before calling `finish()`"
+    # do not add bold() or any other complicated characters as then we can't match this
+    # easily anymore in an html to strip it out
+    return f"please hit {get_shortcut()} to save the notebook in your editor"
+
+
+def get_save_notebook_message_retry() -> str:
+    return f"{get_save_notebook_message()} and re-run finish()"
+
+
+# this code was originally in nbproject by the same authors
+def check_consecutiveness(
+    nb, calling_statement: str = None, silent_success: bool = True
+) -> bool:
+    """Check whether code cells have been executed consecutively.
+
+    Needs to be called in the last code cell of a notebook.
+    Otherwise raises `RuntimeError`.
+
+    Returns cell transitions that violate execution at increments of 1 as a list
+    of tuples.
+
+    Args:
+        nb: Notebook content.
+        calling_statement: The statement that calls this function.
+    """
+    cells = nb.cells
+
+    violations = []
+    prev = 0
+
+    ccount = 0  # need to initialize because notebook might note have code cells
+    # and below, we check if ccount is None
+    for cell in cells:
+        cell_source = "".join(cell["source"])
+        if cell["cell_type"] != "code" or cell_source == "":
+            continue
+
+        if calling_statement is not None and calling_statement in cell_source:
+            continue
+
+        ccount = cell["execution_count"]
+        if ccount is None or prev is None or ccount - prev != 1:
+            violations.append((prev, ccount))
+
+        prev = ccount
+
+    # ignore the very last code cell of the notebook
+    # `check_consecutiveness` is being run during publish if `last_cell`` is True
+    # hence, that cell has ccount is None
+    if ccount is None:
+        violations.pop()
+
+    any_violations = len(violations) > 0
+    if any_violations:
+        logger.warning(f"cells {violations} were not run consecutively")
+    elif not silent_success:
+        logger.success("cell execution numbers increase consecutively")
+
+    return not any_violations
 
 
 def get_shortcut() -> str:
@@ -34,7 +96,7 @@ def save_run_logs(run: Run, save_run: bool = False) -> None:
     if logs_path.exists():
         if run.report is not None:
             logger.important("overwriting run.report")
-        artifact = Artifact(
+        artifact = Artifact(  # type: ignore
             logs_path,
             description=f"log streams of run {run.uid}",
             _branch_code=0,
@@ -64,11 +126,11 @@ def prepare_notebook(
                     if strip_title:
                         lines.pop(i)
                         cell["source"] = "\n".join(lines)
-        # strip resaved finish error if present
+        # strip logging message about saving notebook in editor
         # this is normally the last cell
         if cell["cell_type"] == "code" and ".finish(" in cell["source"]:
             for output in cell["outputs"]:
-                if output.get("ename", None) == "NotebookNotSaved":
+                if "to save the notebook in your editor" in output.get("text", ""):
                     cell["outputs"] = []
                     break
     return None
@@ -97,7 +159,7 @@ def notebook_to_report(notebook_path: Path, output_path: Path) -> None:
     output_path.write_text(html, encoding="utf-8")
 
 
-def notebook_to_script(
+def notebook_to_script(  # type: ignore
     transform: Transform, notebook_path: Path, script_path: Path | None = None
 ) -> None | str:
     import jupytext
@@ -114,7 +176,6 @@ def notebook_to_script(
         script_path.write_text(py_content)
 
 
-# removes NotebookNotSaved error message from notebook html
 def clean_r_notebook_html(file_path: Path) -> tuple[str | None, Path]:
     import re
 
@@ -129,15 +190,15 @@ def clean_r_notebook_html(file_path: Path) -> tuple[str | None, Path]:
         cleaned_content = re.sub(pattern_title, "", cleaned_content)
         cleaned_content = re.sub(pattern_h1, "", cleaned_content)
     # remove error message from content
-    if "NotebookNotSaved" in cleaned_content:
-        orig_error_message = f"NotebookNotSaved: {get_save_notebook_message()}"
+    if "to save the notebook in your editor" in cleaned_content:
+        orig_error_message = f"! {get_save_notebook_message_retry()}"
         # coming up with the regex for this is a bit tricky due to all the
         # escape characters we'd need to insert into the message; hence,
         # we do this with a replace() instead
         cleaned_content = cleaned_content.replace(orig_error_message, "")
-        if "NotebookNotSaved" in cleaned_content:
+        if "to save the notebook in your editor" in cleaned_content:
             orig_error_message = orig_error_message.replace(
-                " `finish()`", "\n`finish()`"
+                " finish()", "\nfinish()"
             )  # RStudio might insert a newline
             cleaned_content = cleaned_content.replace(orig_error_message, "")
     cleaned_path = file_path.parent / (f"{file_path.stem}.cleaned{file_path.suffix}")
@@ -145,6 +206,35 @@ def clean_r_notebook_html(file_path: Path) -> tuple[str | None, Path]:
     return title_text, cleaned_path
 
 
+def check_filepath_recently_saved(filepath: Path, is_finish_retry: bool) -> bool:
+    # the recently_saved_time needs to be very low for the first check
+    # because an accidental save (e.g. via auto-save) might otherwise lead
+    # to upload of an outdated notebook
+    # also see implementation for R notebooks below
+    offset_saved_time = 0.3 if not is_finish_retry else 20
+    for retry in range(30):
+        recently_saved_time = offset_saved_time + retry  # sleep time is 1 sec
+        if get_seconds_since_modified(filepath) > recently_saved_time:
+            if retry == 0:
+                prefix = f"{LEVEL_TO_COLORS[20]}{LEVEL_TO_ICONS[20]}{RESET_COLOR}"
+                print(f"{prefix} {get_save_notebook_message()}", end=" ")
+            elif retry == 9:
+                print(".", end="\n")
+            elif retry == 4:
+                print(". still waiting ", end="")
+            else:
+                print(".", end="")
+            sleep(1)
+        else:
+            if retry > 0:
+                prefix = f"{LEVEL_TO_COLORS[25]}{LEVEL_TO_ICONS[25]}{RESET_COLOR}"
+                print(f" {prefix}")
+            # filepath was recently saved, return True
+            return True
+    # if we arrive here, no save event occured, return False
+    return False
+
+
 def save_context_core(
     *,
     run: Run | None,
@@ -153,6 +243,7 @@ def save_context_core(
     finished_at: bool = False,
     ignore_non_consecutive: bool | None = None,
     from_cli: bool = False,
+    is_retry: bool = False,
 ) -> str | None:
     import lamindb as ln
     from lamindb.models import (
@@ -167,12 +258,29 @@ def save_context_core(
     is_r_notebook = filepath.suffix in {".qmd", ".Rmd"}
     source_code_path = filepath
     report_path: Path | None = None
-    # for notebooks, we need more work
-    if is_ipynb:
+    save_source_code_and_report = True
+    if is_run_from_ipython:  # python notebooks in interactive session
+        import nbproject
+
+        # it might be that the user modifies the title just before ln.finish()
+        if (nbproject_title := nbproject.meta.live.title) != transform.description:
+            transform.description = nbproject_title
+            transform.save()
+        if not ln_setup._TESTING:
+            save_source_code_and_report = check_filepath_recently_saved(
+                filepath, is_retry
+            )
+            if not save_source_code_and_report and not is_retry:
+                logger.warning(get_save_notebook_message_retry())
+                return "retry"
+            elif not save_source_code_and_report:
+                logger.warning(
+                    "the notebook on disk wasn't saved within the last 10 sec"
+                )
+    if is_ipynb:  # could be from CLI outside interactive session
         try:
             import jupytext  # noqa: F401
             from nbproject.dev import (
-                check_consecutiveness,
                 read_notebook,
             )
         except ImportError:
@@ -185,7 +293,9 @@ def save_context_core(
             )
             if not is_consecutive:
                 response = "n"  # ignore_non_consecutive == False
-                if ignore_non_consecutive is None:
+                if ignore_non_consecutive is None:  # only print warning
+                    response = "y"  # we already printed the warning
+                else:  # ask user to confirm
                     response = input(
                         "   Do you still want to proceed with finishing? (y/n) "
                     )
@@ -210,34 +320,45 @@ def save_context_core(
             logger.warning(
                 f"no html report found; to attach one, create an .html export for your {filepath.suffix} file and then run: lamin save {filepath}"
             )
-    if report_path is not None and not from_cli:
-        if get_seconds_since_modified(report_path) > 2 and not ln_setup._TESTING:
-            # this can happen when auto-knitting an html with RStudio
-            raise NotebookNotSaved(get_save_notebook_message())
+    if report_path is not None and is_r_notebook and not from_cli:  # R notebooks
+        # see comment above in check_filepath_recently_saved
+        recently_saved_time = 0.3 if not is_retry else 20
+        if get_seconds_since_modified(report_path) > recently_saved_time:
+            # the automated retry solution of Jupyter notebooks does not work in RStudio because the execution of the notebook cell
+            # seems to block the event loop of the frontend
+            if not is_retry:
+                logger.warning(get_save_notebook_message_retry())
+                return "retry"
+            else:
+                logger.warning(
+                    "the notebook on disk hasn't been saved within the last 20 sec"
+                )
+            save_source_code_and_report = False
     ln.settings.creation.artifact_silence_missing_run_warning = True
-    # track source code
-    hash, _ = hash_file(source_code_path)  # ignore hash_type for now
-    if transform.hash is not None:
-        # check if the hash of the transform source code matches
-        # (for scripts, we already run the same logic in track() - we can deduplicate the call at some point)
-        if hash != transform.hash:
-            response = input(
-                f"You are about to overwrite existing source code (hash '{transform.hash}') for Transform('{transform.uid}')."
-                f" Proceed? (y/n)"
-            )
-            if response == "y":
-                transform.source_code = source_code_path.read_text()
-                transform.hash = hash
+    # save source code
+    if save_source_code_and_report:
+        hash, _ = hash_file(source_code_path)  # ignore hash_type for now
+        if transform.hash is not None:
+            # check if the hash of the transform source code matches
+            # (for scripts, we already run the same logic in track() - we can deduplicate the call at some point)
+            if hash != transform.hash:
+                response = input(
+                    f"You are about to overwrite existing source code (hash '{transform.hash}') for Transform('{transform.uid}')."
+                    f" Proceed? (y/n)"
+                )
+                if response == "y":
+                    transform.source_code = source_code_path.read_text()
+                    transform.hash = hash
+                else:
+                    logger.warning("Please re-run `ln.track()` to make a new version")
+                    return "rerun-the-notebook"
             else:
-                logger.warning("Please re-run `ln.track()` to make a new version")
-                return "rerun-the-notebook"
+                logger.debug("source code is already saved")
         else:
-            logger.debug("source code is already saved")
-    else:
-        transform.source_code = source_code_path.read_text()
-        transform.hash = hash
+            transform.source_code = source_code_path.read_text()
+            transform.hash = hash
 
-    # track environment
+    # track run environment
     if run is not None:
         env_path = ln_setup.settings.cache_dir / f"run_env_pip_{run.uid}.txt"
         if env_path.exists():
@@ -250,7 +371,7 @@ def save_context_core(
                 artifact = ln.Artifact.filter(hash=hash, _branch_code=0).one_or_none()
                 new_env_artifact = artifact is None
                 if new_env_artifact:
-                    artifact = ln.Artifact(
+                    artifact = ln.Artifact(  # type: ignore
                         env_path,
                         description="requirements.txt",
                         _branch_code=0,
@@ -263,50 +384,57 @@ def save_context_core(
 
     # set finished_at
     if finished_at and run is not None:
-        run.finished_at = datetime.now(timezone.utc)
+        if not from_cli:
+            update_finished_at = True
+        else:
+            update_finished_at = run.finished_at is None
+        if update_finished_at:
+            run.finished_at = datetime.now(timezone.utc)
 
     # track logs
     if run is not None and not from_cli and not is_ipynb and not is_r_notebook:
         save_run_logs(run)
 
     # track report and set is_consecutive
-    if run is not None:
-        if report_path is not None:
-            if is_r_notebook:
-                title_text, report_path = clean_r_notebook_html(report_path)
-                if title_text is not None:
-                    transform.description = title_text
-            if run.report_id is not None:
-                hash, _ = hash_file(report_path)  # ignore hash_type for now
-                if hash != run.report.hash:
-                    response = input(
-                        f"You are about to overwrite an existing report (hash '{run.report.hash}') for Run('{run.uid}'). Proceed? (y/n)"
-                    )
-                    if response == "y":
-                        run.report.replace(report_path)
-                        run.report.save(upload=True, print_progress=False)
+    if save_source_code_and_report:
+        if run is not None:
+            if report_path is not None:
+                if is_r_notebook:
+                    title_text, report_path = clean_r_notebook_html(report_path)
+                    if title_text is not None:
+                        transform.description = title_text
+                if run.report_id is not None:
+                    hash, _ = hash_file(report_path)  # ignore hash_type for now
+                    if hash != run.report.hash:
+                        response = input(
+                            f"You are about to overwrite an existing report (hash '{run.report.hash}') for Run('{run.uid}'). Proceed? (y/n)"
+                        )
+                        if response == "y":
+                            run.report.replace(report_path)
+                            run.report.save(upload=True, print_progress=False)
+                        else:
+                            logger.important("keeping old report")
                     else:
-                        logger.important("keeping old report")
+                        logger.important("report is already saved")
                 else:
-                    logger.important("report is already saved")
-            else:
-                report_file = ln.Artifact(
-                    report_path,
-                    description=f"Report of run {run.uid}",
-                    _branch_code=0,  # hidden file
-                    run=False,
+                    report_file = ln.Artifact(  # type: ignore
+                        report_path,
+                        description=f"Report of run {run.uid}",
+                        _branch_code=0,  # hidden file
+                        run=False,
+                    )
+                    report_file.save(upload=True, print_progress=False)
+                    run.report = report_file
+                if is_r_notebook:
+                    # this is the "cleaned" report
+                    report_path.unlink()
+                logger.debug(
+                    f"saved transform.latest_run.report: {transform.latest_run.report}"
                 )
-                report_file.save(upload=True, print_progress=False)
-                run.report = report_file
-            if is_r_notebook:
-                # this is the "cleaned" report
-                report_path.unlink()
-            logger.debug(
-                f"saved transform.latest_run.report: {transform.latest_run.report}"
-            )
-        run._is_consecutive = is_consecutive
+            run._is_consecutive = is_consecutive
 
-        # save both run & transform records if we arrive here
+    # save both run & transform records if we arrive here
+    if run is not None:
         run.save()
     transform.save()
 
@@ -318,19 +446,32 @@ def save_context_core(
         hours = seconds // 3600
         minutes = (seconds % 3600) // 60
         secs = seconds % 60
-        formatted_run_time = f"{days}d {hours}h {minutes}m {secs}s"
+        formatted_run_time = (
+            f"{days}d"
+            if days != 0
+            else "" + f"{hours}h"
+            if hours != 0
+            else "" + f"{minutes}m"
+            if minutes != 0
+            else "" + f"{secs}s"
+        )
 
         logger.important(
             f"finished Run('{run.uid[:8]}') after {formatted_run_time} at {format_field_value(run.finished_at)}"
         )
     if ln_setup.settings.instance.is_on_hub:
-        identifier = ln_setup.settings.instance.slug
-        logger.important(
-            f"go to: https://lamin.ai/{identifier}/transform/{transform.uid}"
-        )
-        if not from_cli:
+        instance_slug = ln_setup.settings.instance.slug
+        if save_source_code_and_report:
+            logger.important(
+                f"go to: https://lamin.ai/{instance_slug}/transform/{transform.uid}"
+            )
+        if not from_cli and save_source_code_and_report:
             thing = "notebook" if (is_ipynb or is_r_notebook) else "script"
             logger.important(
-                f"if you want to update your {thing} without re-running it, use `lamin save {filepath}`"
+                f"to update your {thing} from the CLI, run: lamin save {filepath}"
             )
+    if not save_source_code_and_report:
+        logger.warning(
+            f"did *not* save source code and report -- to do so, run: lamin save {filepath}"
+        )
     return None