From 848bc82a9cdc56c64fa288554dd84b4be323f5e6 Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Fri, 6 Nov 2026 21:52:01 +0000 Subject: [PATCH] chore(scripts): add reusable notebook_setup sparse-checkout helper Add a reusable, standard-library bootstrap utility in scripts/notebook-setup/ for tutorials and sample notebooks that depend on sibling Python modules when opened standalone in Google Colab or Colab Enterprise. - Sparse-clones only the target sample subdirectory into a managed cache folder, self-heals interrupted non-git directories, fetches and checks out FETCH_HEAD on persistent runtimes, installs requirements.txt, and reloads imported sibling sample modules. - Preserves existing local developer and CI repository checkouts without running git fetch or git checkout -f. - Includes unit tests in scripts/notebook-setup/tests/notebook_setup_test.py. --- scripts/notebook-setup/README.md | 31 +++ scripts/notebook-setup/__init__.py | 13 + scripts/notebook-setup/notebook_setup.py | 210 ++++++++++++++++ scripts/notebook-setup/noxfile_config.py | 21 ++ scripts/notebook-setup/requirements-test.txt | 1 + scripts/notebook-setup/requirements.txt | 1 + scripts/notebook-setup/tests/__init__.py | 13 + .../tests/notebook_setup_test.py | 225 ++++++++++++++++++ 8 files changed, 515 insertions(+) create mode 100644 scripts/notebook-setup/README.md create mode 100644 scripts/notebook-setup/__init__.py create mode 100644 scripts/notebook-setup/notebook_setup.py create mode 100644 scripts/notebook-setup/noxfile_config.py create mode 100644 scripts/notebook-setup/requirements-test.txt create mode 100644 scripts/notebook-setup/requirements.txt create mode 100644 scripts/notebook-setup/tests/__init__.py create mode 100644 scripts/notebook-setup/tests/notebook_setup_test.py diff --git a/scripts/notebook-setup/README.md b/scripts/notebook-setup/README.md new file mode 100644 index 00000000000..c78f944e8f7 --- /dev/null +++ b/scripts/notebook-setup/README.md @@ -0,0 +1,31 @@ +# Notebook Setup Helper (`notebook_setup.py`) + +Reusable standard-library bootstrap utility for notebooks in `python-docs-samples` that depend on sibling Python modules and `requirements.txt` files when opened standalone in Google Colab or Colab Enterprise. + +## Usage in a Notebook + +```python +# @title Download sample files and install packages (~30 seconds) +from pathlib import Path +import sys +import urllib.request + +BRANCH = "main" +SETUP_PATH = Path("notebook_setup.py") +for parent in (Path.cwd(), *Path.cwd().parents): + candidate = parent / "scripts/notebook-setup/notebook_setup.py" + if candidate.is_file(): + SETUP_PATH = candidate + break +else: + if not SETUP_PATH.is_file(): + urllib.request.urlretrieve( + f"https://raw.githubusercontent.com/GoogleCloudPlatform/python-docs-samples/{BRANCH}/scripts/notebook-setup/notebook_setup.py", + SETUP_PATH, + ) + +sys.path.insert(0, str(SETUP_PATH.parent.resolve())) +import notebook_setup + +notebook_setup.bootstrap("people-and-planet-ai/weather-forecasting", branch=BRANCH) +``` diff --git a/scripts/notebook-setup/__init__.py b/scripts/notebook-setup/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/scripts/notebook-setup/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/scripts/notebook-setup/notebook_setup.py b/scripts/notebook-setup/notebook_setup.py new file mode 100644 index 00000000000..1cd9b92c848 --- /dev/null +++ b/scripts/notebook-setup/notebook_setup.py @@ -0,0 +1,210 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reusable sparse-checkout and package setup helper for sample notebooks.""" + +from __future__ import annotations + +from collections.abc import Sequence +import importlib +import os +from pathlib import Path +import shutil +import subprocess +import sys + +DEFAULT_REPO_URL = "https://github.com/GoogleCloudPlatform/python-docs-samples.git" +_MANAGED_MARKER = "notebook_setup_managed" + + +def _find_git_root(start: Path) -> Path | None: + """Walks upward from `start` to find an enclosing `.git` directory.""" + current = start.resolve() + for candidate in (current, *current.parents): + if (candidate / ".git").exists(): + return candidate + return None + + +def _is_managed_clone(repo_root: Path) -> bool: + """Returns True if `repo_root` was created and managed by `notebook_setup`.""" + return (repo_root / ".git" / _MANAGED_MARKER).is_file() + + +def _reload_sample_modules( + sample_dir: Path, + extra_modules: Sequence[str] = (), +) -> None: + """Reloads already-imported sibling Python modules in `sample_dir`.""" + importlib.invalidate_caches() + local_stems = { + py_file.stem + for py_file in sample_dir.glob("*.py") + if py_file.stem != "__init__" and not py_file.stem.startswith("test_") + } + local_stems.update(extra_modules) + for mod_name in sorted(local_stems): + module = sys.modules.get(mod_name) + if module is not None: + importlib.reload(module) + + +def bootstrap( + sample_subpath: str, + *, + branch: str = "main", + repo_url: str = DEFAULT_REPO_URL, + clone_dir: str | None = None, + install_requirements: bool = True, + reload_modules: Sequence[str] = (), +) -> Path: + """Sparse-clones or updates `sample_subpath` and installs its dependencies. + + When executed inside an existing local developer or CI checkout of the + repository, `bootstrap` leaves local working-tree files untouched (no + `git fetch` or `git checkout -f`), switches into `sample_subpath`, updates + `sys.path`, and optionally installs `requirements.txt`. + + When executed in an ephemeral or persistent notebook environment (such as + Google Colab or Colab Enterprise) where only the `.ipynb` file was imported, + `bootstrap` creates or updates a managed sparse checkout, self-heals + interrupted partial directories, checks out the latest `FETCH_HEAD` for + `branch`, installs `requirements.txt`, and reloads any previously imported + sibling sample modules. + + Args: + sample_subpath: Repository-relative path to the sample directory (for + example, `'people-and-planet-ai/weather-forecasting'`). + branch: Git branch or tag to fetch when managing a notebook clone. + repo_url: Git remote URL to sparse-clone when not in a local checkout. + clone_dir: Optional directory name for the managed sparse clone. + Defaults to `'samples-'`. + install_requirements: Whether to install `requirements.txt` with `pip`. + reload_modules: Optional extra module names to reload if already in + `sys.modules`. + + Returns: + Resolved `Path` to the active sample directory. + """ + normalized_subpath = sample_subpath.strip().strip("/") + if not normalized_subpath or normalized_subpath in {".", ".."}: + raise ValueError("sample_subpath must be a non-empty relative path.") + + cwd = Path.cwd().resolve() + enclosing_git_root = _find_git_root(cwd) + + if enclosing_git_root is not None and not _is_managed_clone(enclosing_git_root): + sample_dir = (enclosing_git_root / normalized_subpath).resolve() + if not sample_dir.is_dir(): + raise FileNotFoundError( + f"Sample directory '{normalized_subpath}' not found inside " + f"local repository '{enclosing_git_root}'." + ) + else: + if enclosing_git_root is not None and _is_managed_clone(enclosing_git_root): + repo_root = enclosing_git_root + else: + leaf_name = Path(normalized_subpath).name + target_dir_name = clone_dir or f"samples-{leaf_name}" + repo_root = (cwd / target_dir_name).resolve() + + if repo_root.exists() and not (repo_root / ".git").exists(): + if repo_root.is_dir(): + shutil.rmtree(repo_root) + else: + repo_root.unlink() + + if not repo_root.exists(): + subprocess.run( + [ + "git", + "clone", + "--filter=blob:none", + "--no-checkout", + "--depth", + "1", + "--sparse", + "-b", + branch, + repo_url, + str(repo_root), + ], + check=True, + ) + (repo_root / ".git" / _MANAGED_MARKER).write_text( + "managed-by-notebook-setup\n", + encoding="utf-8", + ) + + subprocess.run( + [ + "git", + "-C", + str(repo_root), + "sparse-checkout", + "set", + normalized_subpath, + ], + check=True, + ) + subprocess.run( + [ + "git", + "-C", + str(repo_root), + "fetch", + "--depth", + "1", + "origin", + branch, + ], + check=True, + ) + subprocess.run( + [ + "git", + "-C", + str(repo_root), + "-c", + "advice.detachedHead=false", + "checkout", + "-f", + "FETCH_HEAD", + ], + check=True, + ) + sample_dir = (repo_root / normalized_subpath).resolve() + + os.chdir(sample_dir) + sample_dir_str = str(sample_dir) + if sample_dir_str not in sys.path: + sys.path.insert(0, sample_dir_str) + + req_file = sample_dir / "requirements.txt" + if install_requirements and req_file.is_file(): + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--quiet", + "-r", + str(req_file), + ], + check=True, + ) + + _reload_sample_modules(sample_dir, extra_modules=reload_modules) + return sample_dir diff --git a/scripts/notebook-setup/noxfile_config.py b/scripts/notebook-setup/noxfile_config.py new file mode 100644 index 00000000000..8e6ff7143a9 --- /dev/null +++ b/scripts/notebook-setup/noxfile_config.py @@ -0,0 +1,21 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +TEST_CONFIG_OVERRIDE = { + "ignored_versions": ["3.8", "3.9", "3.11", "3.12", "3.13"], + "enforce_type_hints": True, + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + "pip_version_override": None, + "envs": {}, +} diff --git a/scripts/notebook-setup/requirements-test.txt b/scripts/notebook-setup/requirements-test.txt new file mode 100644 index 00000000000..e144ed090b5 --- /dev/null +++ b/scripts/notebook-setup/requirements-test.txt @@ -0,0 +1 @@ +pytest>=8.0.0 diff --git a/scripts/notebook-setup/requirements.txt b/scripts/notebook-setup/requirements.txt new file mode 100644 index 00000000000..7d0a78e3f4d --- /dev/null +++ b/scripts/notebook-setup/requirements.txt @@ -0,0 +1 @@ +# Runtime uses only the Python standard library; kept for CI test discovery. diff --git a/scripts/notebook-setup/tests/__init__.py b/scripts/notebook-setup/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/scripts/notebook-setup/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/scripts/notebook-setup/tests/notebook_setup_test.py b/scripts/notebook-setup/tests/notebook_setup_test.py new file mode 100644 index 00000000000..3ef85dafae5 --- /dev/null +++ b/scripts/notebook-setup/tests/notebook_setup_test.py @@ -0,0 +1,225 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the reusable notebook_setup helper.""" + +from __future__ import annotations + +import importlib +import importlib.util +from pathlib import Path +import subprocess +import sys +from types import ModuleType +from unittest import mock + +import pytest + +_SPEC = importlib.util.spec_from_file_location( + "notebook_setup", + Path(__file__).resolve().parents[1] / "notebook_setup.py", +) +assert _SPEC is not None and _SPEC.loader is not None +notebook_setup = importlib.util.module_from_spec(_SPEC) +sys.modules["notebook_setup"] = notebook_setup +_SPEC.loader.exec_module(notebook_setup) + +SAMPLE_SUBPATH = "people-and-planet-ai/weather-forecasting" + + +def test_bootstrap_fresh_colab_sparse_checkout( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", list(sys.path)) + + def fake_run( + cmd: list[str], + *, + check: bool = False, + ) -> subprocess.CompletedProcess[str]: + assert check is True + if cmd[:2] == ["git", "clone"]: + repo_root = Path(cmd[-1]) + (repo_root / ".git").mkdir(parents=True, exist_ok=True) + elif "checkout" in cmd and "FETCH_HEAD" in cmd: + repo_root = Path(cmd[2]) + sample_dir = repo_root / SAMPLE_SUBPATH + sample_dir.mkdir(parents=True, exist_ok=True) + (sample_dir / "requirements.txt").write_text( + "pytest>=8.0.0\n", encoding="utf-8" + ) + (sample_dir / "visualize.py").write_text("VALUE = 1\n", encoding="utf-8") + return subprocess.CompletedProcess(cmd, 0) + + with mock.patch.object( + notebook_setup.subprocess, "run", side_effect=fake_run, autospec=True + ) as mock_run: + resolved = notebook_setup.bootstrap( + SAMPLE_SUBPATH, + branch="feat/test-branch", + ) + + expected_repo = tmp_path / "samples-weather-forecasting" + expected_sample_dir = expected_repo / SAMPLE_SUBPATH + assert resolved == expected_sample_dir + assert Path.cwd() == expected_sample_dir + assert sys.path[0] == str(expected_sample_dir) + assert (expected_repo / ".git" / "notebook_setup_managed").is_file() + # 4 git commands (clone, sparse-checkout, fetch, checkout) + 1 pip install + assert mock_run.call_count == 5 + pip_cmd = mock_run.call_args_list[-1].args[0] + assert pip_cmd[:4] == [sys.executable, "-m", "pip", "install"] + assert "--upgrade" not in pip_cmd + + +def test_bootstrap_heals_broken_non_git_cache_and_reloads_module( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", list(sys.path)) + + broken_dir = tmp_path / "samples-weather-forecasting" + broken_dir.mkdir() + stale_file = broken_dir / "orphan.txt" + stale_file.write_text("partial", encoding="utf-8") + + fake_mod = ModuleType("visualize") + monkeypatch.setitem(sys.modules, "visualize", fake_mod) + monkeypatch.setitem(sys.modules, "disabled_mod", None) + + def fake_run( + cmd: list[str], + *, + check: bool = False, + ) -> subprocess.CompletedProcess[str]: + assert check is True + if cmd[:2] == ["git", "clone"]: + assert not stale_file.exists() + repo_root = Path(cmd[-1]) + (repo_root / ".git").mkdir(parents=True, exist_ok=True) + elif "checkout" in cmd and "FETCH_HEAD" in cmd: + repo_root = Path(cmd[2]) + sample_dir = repo_root / SAMPLE_SUBPATH + sample_dir.mkdir(parents=True, exist_ok=True) + (sample_dir / "visualize.py").write_text("VALUE = 2\n", encoding="utf-8") + (sample_dir / "disabled_mod.py").write_text("VALUE = 3\n", encoding="utf-8") + return subprocess.CompletedProcess(cmd, 0) + + with ( + mock.patch.object( + notebook_setup.subprocess, "run", side_effect=fake_run, autospec=True + ), + mock.patch.object(importlib, "reload", autospec=True) as mock_reload, + ): + notebook_setup.bootstrap(SAMPLE_SUBPATH, install_requirements=False) + + mock_reload.assert_called_once_with(fake_mod) + + +def test_bootstrap_unlinks_conflicting_file_at_repo_root( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", list(sys.path)) + + conflicting_file = tmp_path / "samples-weather-forecasting" + conflicting_file.write_text("accidental file", encoding="utf-8") + + def fake_run( + cmd: list[str], + *, + check: bool = False, + ) -> subprocess.CompletedProcess[str]: + assert check is True + if cmd[:2] == ["git", "clone"]: + repo_root = Path(cmd[-1]) + assert not repo_root.is_file() + (repo_root / ".git").mkdir(parents=True, exist_ok=True) + elif "checkout" in cmd and "FETCH_HEAD" in cmd: + repo_root = Path(cmd[2]) + (repo_root / SAMPLE_SUBPATH).mkdir(parents=True, exist_ok=True) + return subprocess.CompletedProcess(cmd, 0) + + with mock.patch.object( + notebook_setup.subprocess, "run", side_effect=fake_run, autospec=True + ): + resolved = notebook_setup.bootstrap(SAMPLE_SUBPATH, install_requirements=False) + + assert resolved.is_dir() + + +def test_bootstrap_updates_when_already_inside_managed_clone( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + repo_root = tmp_path / "samples-weather-forecasting" + (repo_root / ".git").mkdir(parents=True) + (repo_root / ".git" / "notebook_setup_managed").write_text( + "managed\n", encoding="utf-8" + ) + sample_dir = repo_root / SAMPLE_SUBPATH + sample_dir.mkdir(parents=True) + monkeypatch.chdir(sample_dir) + monkeypatch.setattr(sys, "path", list(sys.path)) + + with mock.patch.object(notebook_setup.subprocess, "run", autospec=True) as mock_run: + resolved = notebook_setup.bootstrap( + SAMPLE_SUBPATH, + branch="main", + install_requirements=False, + ) + + assert resolved == sample_dir + # Should run sparse-checkout, fetch, and checkout (no git clone) + cmds = [call.args[0] for call in mock_run.call_args_list] + assert len(cmds) == 3 + assert cmds[0][:2] == ["git", "-C"] + assert "sparse-checkout" in cmds[0] + assert "fetch" in cmds[1] + assert "checkout" in cmds[2] + + +def test_bootstrap_preserves_local_developer_checkout_without_git_reset( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + local_repo = tmp_path / "python-docs-samples" + (local_repo / ".git").mkdir(parents=True) + sample_dir = local_repo / SAMPLE_SUBPATH + sample_dir.mkdir(parents=True) + local_edit = sample_dir / "uncommitted.py" + local_edit.write_text("LOCAL = True\n", encoding="utf-8") + + monkeypatch.chdir(local_repo) + monkeypatch.setattr(sys, "path", list(sys.path)) + + with mock.patch.object(notebook_setup.subprocess, "run", autospec=True) as mock_run: + resolved = notebook_setup.bootstrap( + SAMPLE_SUBPATH, + install_requirements=False, + ) + + assert resolved == sample_dir + assert Path.cwd() == sample_dir + assert local_edit.read_text(encoding="utf-8") == "LOCAL = True\n" + mock_run.assert_not_called() + + +def test_bootstrap_rejects_empty_subpath() -> None: + with pytest.raises(ValueError, match="non-empty"): + notebook_setup.bootstrap(" / ")