From fff0dd1f17c76691550267bfc92a73a4840b8b19 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 | 36 ++ scripts/notebook-setup/__init__.py | 13 + scripts/notebook-setup/notebook_setup.py | 375 ++++++++++++++++++ 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 | 319 +++++++++++++++ 8 files changed, 779 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..104df43573e --- /dev/null +++ b/scripts/notebook-setup/README.md @@ -0,0 +1,36 @@ +# Notebook Setup Helper (`notebook_setup.py`) + +`notebook_setup.py` provides a 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 tempfile +import urllib.request + +BRANCH = "main" +SETUP_PATH = Path(tempfile.gettempdir()) / "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(): + url = ( + "https://raw.githubusercontent.com/GoogleCloudPlatform/" + f"python-docs-samples/{BRANCH}/scripts/notebook-setup/notebook_setup.py" + ) + with urllib.request.urlopen(url, timeout=30) as response: + SETUP_PATH.write_bytes(response.read()) + +setup_dir = str(SETUP_PATH.parent.resolve()) +if setup_dir not in sys.path: + sys.path.insert(0, setup_dir) +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..3e7894c90a0 --- /dev/null +++ b/scripts/notebook-setup/notebook_setup.py @@ -0,0 +1,375 @@ +# 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, PurePosixPath +import shutil +import subprocess +import sys + +DEFAULT_REPO_URL = "https://github.com/GoogleCloudPlatform/python-docs-samples.git" +_MANAGED_MARKER = "notebook_setup_managed" +_NETWORK_TIMEOUT_SECONDS = 300.0 +_LOCAL_GIT_TIMEOUT_SECONDS = 30.0 + + +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 _validate_inputs( + sample_subpath: str, + branch: str, + repo_url: str, + clone_dir: str | None, +) -> tuple[str, str, str]: + """Validates path and git arguments and returns normalized values.""" + raw_subpath = sample_subpath.strip().strip("/") + parts = PurePosixPath(raw_subpath).parts + if ( + not parts + or "\\" in sample_subpath + or any(part in {"", ".", ".."} for part in parts) + or parts[0].startswith("-") + ): + raise ValueError( + f"sample_subpath {sample_subpath!r} must be a non-empty relative path " + "without '..' segments or leading '-'." + ) + + clean_branch = branch.strip() + if ( + not clean_branch + or clean_branch.startswith("-") + or any(ch.isspace() for ch in clean_branch) + ): + raise ValueError(f"Invalid git branch or tag {branch!r}.") + + clean_repo_url = repo_url.strip() + if ( + not clean_repo_url + or clean_repo_url.startswith("-") + or any(ch.isspace() for ch in clean_repo_url) + ): + raise ValueError(f"Invalid repo_url {repo_url!r}.") + + if clone_dir is not None: + raw_clone = clone_dir.strip() + clone_parts = PurePosixPath(raw_clone).parts + if ( + len(clone_parts) != 1 + or "\\" in raw_clone + or raw_clone.startswith("/") + or clone_parts[0] in {"", ".", ".."} + or clone_parts[0].startswith("-") + ): + raise ValueError( + f"clone_dir {clone_dir!r} must be a single relative directory name." + ) + + return "/".join(parts), clean_branch, clean_repo_url + + +def _run_command(cmd: Sequence[str], *, timeout: float) -> None: + """Executes a subprocess non-interactively and surfaces stderr on failure.""" + env = { + **os.environ, + "GIT_TERMINAL_PROMPT": "0", + "PIP_NO_INPUT": "1", + } + cmd_list = list(cmd) + try: + subprocess.run( + cmd_list, + check=True, + capture_output=True, + text=True, + timeout=timeout, + env=env, + ) + except subprocess.CalledProcessError as err: + detail = (err.stderr or err.stdout or "").strip() + suffix = f": {detail}" if detail else "." + raise RuntimeError( + f"Command {cmd_list!r} failed with exit code {err.returncode}{suffix}" + ) from err + + +def _sync_managed_clone( + repo_root: Path, + normalized_subpath: str, + branch: str, + repo_url: str, +) -> Path: + """Creates or updates a managed sparse checkout and returns `sample_dir`.""" + if (repo_root / ".git").exists() and not _is_managed_clone(repo_root): + raise RuntimeError( + f"Refusing to modify existing unmanaged git repository at '{repo_root}'. " + "Pass a distinct clone_dir or run inside that checkout." + ) + + if repo_root.is_symlink() or ( + repo_root.exists() and not (repo_root / ".git").exists() + ): + if not repo_root.is_symlink() and repo_root.is_dir(): + shutil.rmtree(repo_root) + else: + repo_root.unlink() + + is_fresh_clone = not repo_root.exists() + if is_fresh_clone: + _run_command( + [ + "git", + "clone", + "--filter=blob:none", + "--no-checkout", + "--depth", + "1", + "--sparse", + "-b", + branch, + "--", + repo_url, + str(repo_root), + ], + timeout=_NETWORK_TIMEOUT_SECONDS, + ) + (repo_root / ".git" / _MANAGED_MARKER).write_text( + "managed-by-notebook-setup\n", + encoding="utf-8", + ) + + _run_command( + [ + "git", + "-C", + str(repo_root), + "sparse-checkout", + "set", + "--", + normalized_subpath, + ], + timeout=_LOCAL_GIT_TIMEOUT_SECONDS, + ) + if is_fresh_clone: + _run_command( + [ + "git", + "-C", + str(repo_root), + "-c", + "advice.detachedHead=false", + "checkout", + "-f", + "HEAD", + ], + timeout=_NETWORK_TIMEOUT_SECONDS, + ) + else: + _run_command( + [ + "git", + "-C", + str(repo_root), + "fetch", + "--depth", + "1", + "--", + "origin", + branch, + ], + timeout=_NETWORK_TIMEOUT_SECONDS, + ) + _run_command( + [ + "git", + "-C", + str(repo_root), + "-c", + "advice.detachedHead=false", + "checkout", + "-f", + "FETCH_HEAD", + ], + timeout=_NETWORK_TIMEOUT_SECONDS, + ) + + sample_dir = (repo_root / normalized_subpath).resolve() + if not sample_dir.is_relative_to(repo_root) or not sample_dir.is_dir(): + raise FileNotFoundError( + f"Sample directory '{normalized_subpath}' not found in '{repo_root}' " + f"(branch '{branch}')." + ) + return sample_dir + + +def _is_reloadable_sample_file(py_file: Path) -> bool: + """Returns True if `py_file` is a sample module rather than test/CI config.""" + stem = py_file.stem + return ( + stem + not in { + "__init__", + "conftest", + "notebook_setup", + "noxfile", + "noxfile_config", + } + and not stem.startswith("test_") + and not stem.endswith("_test") + ) + + +def _reload_sample_modules( + sample_dir: Path, + extra_modules: Sequence[str] = (), +) -> None: + """Reloads already-imported sibling Python modules originating in `sample_dir`.""" + importlib.invalidate_caches() + resolved_dir = sample_dir.resolve() + explicit = set(extra_modules) + candidates = { + py_file.stem + for py_file in resolved_dir.glob("*.py") + if _is_reloadable_sample_file(py_file) + } | explicit + for mod_name in sorted(candidates): + module = sys.modules.get(mod_name) + if module is None: + continue + mod_file = getattr(module, "__file__", None) + if mod_name in explicit or ( + mod_file is not None + and Path(mod_file).resolve().is_relative_to(resolved_dir) + ): + 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 target `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 single relative 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. + + Raises: + ValueError: If `sample_subpath`, `branch`, `repo_url`, or `clone_dir` + is empty or contains unsafe path/flag segments. + FileNotFoundError: If `sample_subpath` does not exist in the repository. + RuntimeError: If `repo_root` is an unmanaged git checkout or if a `git` + or `pip` command fails. + subprocess.TimeoutExpired: If a `git` or `pip` subprocess exceeds its + timeout. + """ + normalized_subpath, clean_branch, clean_repo_url = _validate_inputs( + sample_subpath, branch, repo_url, clone_dir + ) + + 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_relative_to(enclosing_git_root) or 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 = PurePosixPath(normalized_subpath).name + target_dir_name = ( + clone_dir.strip() if clone_dir is not None else f"samples-{leaf_name}" + ) + repo_root = cwd / target_dir_name + + sample_dir = _sync_managed_clone( + repo_root=repo_root, + normalized_subpath=normalized_subpath, + branch=clean_branch, + repo_url=clean_repo_url, + ) + + os.chdir(sample_dir) + sample_dir_str = str(sample_dir) + if sample_dir_str in sys.path: + sys.path.remove(sample_dir_str) + sys.path.insert(0, sample_dir_str) + + req_file = sample_dir / "requirements.txt" + if install_requirements and req_file.is_file(): + _run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "--quiet", + "-r", + str(req_file), + ], + timeout=_NETWORK_TIMEOUT_SECONDS, + ) + + _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..1ca783926d1 --- /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.10", "3.11", "3.13", "3.14"], + "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..7345f250f23 --- /dev/null +++ b/scripts/notebook-setup/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.1.1 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..dee777ae25a --- /dev/null +++ b/scripts/notebook-setup/tests/notebook_setup_test.py @@ -0,0 +1,319 @@ +# 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 + +from collections.abc import Callable, Mapping +import importlib +import importlib.util +from pathlib import Path +import subprocess +import sys +import types +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 _make_fake_git_runner( + on_checkout: Callable[[Path], None] | None = None, +) -> Callable[..., subprocess.CompletedProcess[str]]: + """Builds a `subprocess.run` fake that materializes sparse git directories.""" + + def _fake_run( + cmd: list[str], + *, + check: bool = False, + capture_output: bool = False, + text: bool = False, + timeout: float | None = None, + env: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + assert check is True + assert capture_output is True + assert text is True + assert timeout is not None + assert env is not None and env.get("GIT_TERMINAL_PROMPT") == "0" + if cmd[:2] == ["git", "clone"]: + repo_root = Path(cmd[-1]) + (repo_root / ".git").mkdir(parents=True, exist_ok=True) + elif "checkout" in cmd and ("HEAD" in cmd or "FETCH_HEAD" in cmd): + repo_root = Path(cmd[2]) + sample_dir = repo_root / SAMPLE_SUBPATH + sample_dir.mkdir(parents=True, exist_ok=True) + if on_checkout is not None: + on_checkout(sample_dir) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + return _fake_run + + +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 populate_files(sample_dir: Path) -> None: + (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") + + with mock.patch.object( + notebook_setup.subprocess, + "run", + side_effect=_make_fake_git_runner(populate_files), + 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() + + # 3 git commands on cold clone (clone, sparse-checkout, checkout HEAD) + 1 pip + cmds = [call.args[0] for call in mock_run.call_args_list] + assert len(cmds) == 4, f"got {len(cmds)} commands ({cmds!r}), want 4" + assert "feat/test-branch" in cmds[0] + assert "--" in cmds[0] + assert cmds[1][-2:] == ["--", SAMPLE_SUBPATH] + assert cmds[2][-2:] == ["-f", "HEAD"] + assert all("fetch" not in cmd for cmd in cmds) + assert cmds[3][:4] == [sys.executable, "-m", "pip", "install"] + assert "--upgrade" not in cmds[3] + + +def test_bootstrap_heals_broken_non_git_cache_and_isolates_reloads( + 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") + + expected_sample_dir = broken_dir / SAMPLE_SUBPATH + fake_mod = ModuleType("visualize") + fake_mod.__file__ = str(expected_sample_dir / "visualize.py") + monkeypatch.setitem(sys.modules, "visualize", fake_mod) + monkeypatch.setitem(sys.modules, "disabled_mod", None) + assert sys.modules.get("types") is types + + def populate_files(sample_dir: Path) -> None: + assert not stale_file.exists() + (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") + # Collides by stem with Python stdlib `types`, which must NOT be reloaded. + (sample_dir / "types.py").write_text("SHADOW = True\n", encoding="utf-8") + (sample_dir / "visualize_test.py").write_text("TEST = True\n", encoding="utf-8") + + with ( + mock.patch.object( + notebook_setup.subprocess, + "run", + side_effect=_make_fake_git_runner(populate_files), + 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_or_symlink_at_repo_root( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(sys, "path", list(sys.path)) + + target_dir = tmp_path / "external-target" + target_dir.mkdir() + keep_file = target_dir / "keep.txt" + keep_file.write_text("must survive\n", encoding="utf-8") + + symlink_root = tmp_path / "samples-weather-forecasting" + symlink_root.symlink_to(target_dir, target_is_directory=True) + + with mock.patch.object( + notebook_setup.subprocess, + "run", + side_effect=_make_fake_git_runner(), + autospec=True, + ): + resolved = notebook_setup.bootstrap(SAMPLE_SUBPATH, install_requirements=False) + + assert resolved.is_dir() + assert keep_file.read_text(encoding="utf-8") == "must survive\n" + + +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", ["/other/path", str(sample_dir)]) + + with mock.patch.object( + notebook_setup.subprocess, + "run", + return_value=subprocess.CompletedProcess(["git"], 0, stdout="", stderr=""), + autospec=True, + ) as mock_run: + resolved = notebook_setup.bootstrap( + SAMPLE_SUBPATH, + branch="main", + install_requirements=False, + ) + + assert resolved == sample_dir + assert sys.path[0] == str(sample_dir) + # Warm clone runs sparse-checkout, fetch, and checkout FETCH_HEAD (no clone) + cmds = [call.args[0] for call in mock_run.call_args_list] + assert len(cmds) == 3, f"got {len(cmds)} commands ({cmds!r}), want 3" + assert "sparse-checkout" in cmds[0] + assert "fetch" in cmds[1] and "main" in cmds[1] + assert cmds[2][-2:] == ["-f", "FETCH_HEAD"] + + +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_refuses_unmanaged_git_repo_at_clone_dir( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + unmanaged_repo = tmp_path / "python-docs-samples" + (unmanaged_repo / ".git").mkdir(parents=True) + monkeypatch.chdir(tmp_path) + + with ( + mock.patch.object(notebook_setup.subprocess, "run", autospec=True) as mock_run, + pytest.raises(RuntimeError, match="unmanaged git repository"), + ): + notebook_setup.bootstrap( + SAMPLE_SUBPATH, + clone_dir="python-docs-samples", + install_requirements=False, + ) + + mock_run.assert_not_called() + + +def test_bootstrap_raises_file_not_found_for_missing_sample_dir( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + local_repo = tmp_path / "python-docs-samples" + (local_repo / ".git").mkdir(parents=True) + monkeypatch.chdir(local_repo) + + with pytest.raises(FileNotFoundError, match="not found"): + notebook_setup.bootstrap("nonexistent/sample", install_requirements=False) + + +def test_bootstrap_surfaces_subprocess_stderr_on_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + err = subprocess.CalledProcessError( + 128, + ["git", "clone"], + output="", + stderr="fatal: Remote branch missing-branch not found in upstream origin\n", + ) + with ( + mock.patch.object( + notebook_setup.subprocess, "run", side_effect=err, autospec=True + ), + pytest.raises(RuntimeError, match="Remote branch missing-branch not found"), + ): + notebook_setup.bootstrap(SAMPLE_SUBPATH, branch="missing-branch") + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"sample_subpath": " / "}, "non-empty"), + ({"sample_subpath": "foo/../../outside"}, "relative path"), + ({"sample_subpath": "-upload-pack=bad"}, "leading '-'"), + ({"sample_subpath": SAMPLE_SUBPATH, "clone_dir": "."}, "single relative"), + ({"sample_subpath": SAMPLE_SUBPATH, "clone_dir": "../sibling"}, "single"), + ({"sample_subpath": SAMPLE_SUBPATH, "branch": "--upload-pack=x"}, "branch"), + ({"sample_subpath": SAMPLE_SUBPATH, "repo_url": "--upload-pack=x"}, "repo_url"), + ], +) +def test_bootstrap_rejects_unsafe_inputs( + kwargs: dict[str, str], + match: str, +) -> None: + with pytest.raises(ValueError, match=match): + notebook_setup.bootstrap(**kwargs)