googlecloudplatform--generative-ai/scripts/notebook_processor.py

# pylint: skip-file
# type: ignore
# -*- coding: utf-8 -*-
#
# Copyright 2024 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
#
#     https://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.

import os
import sys
import urllib.parse

import nbformat

DEFAULT_RUFF_LINE_LENGTH = 88

# --- Notebook Link Constants ---
LINK_PREFIXES = {
    "colab_link": "https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/",
    "colab_enterprise_link": "https://console.cloud.google.com/agent-platform/colab/import/",
    "github_link": "https://github.com/GoogleCloudPlatform/generative-ai/blob/main/",
    "workbench_link": "https://console.cloud.google.com/agent-platform/workbench/instances?download_url=",
    "bigquery_studio_link": "https://console.cloud.google.com/bigquery/import?url=",
    "linkedin_link": "https://www.linkedin.com/sharing/share-offsite/?url=",
    "bluesky_link": "https://bsky.app/intent/compose?text=",
    "twitter_link": "https://twitter.com/intent/tweet?url=",
    "reddit_link": "https://reddit.com/submit?url=",
    "facebook_link": "https://www.facebook.com/sharer/sharer.php?u=",
}

GITHUB_URL_PREFIX = LINK_PREFIXES["github_link"]
RAW_URL_PREFIX = (
    "https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/"
)


def fix_markdown_links(
    cell_source: str, relative_notebook_path: str
) -> tuple[str, bool]:
    """Fixes links in a markdown cell and returns the updated source."""
    new_lines = []
    changes_made = False

    encoded_url = urllib.parse.quote(f"{GITHUB_URL_PREFIX}{relative_notebook_path}")

    for line in cell_source.splitlines():
        for key, prefix in LINK_PREFIXES.items():
            if prefix not in line or "**NOTE:**" in line:
                continue

            start_index = line.find(prefix) + len(prefix)
            end_index = line.find(".ipynb", start_index) + len(".ipynb")
            correct_link = ""

            if key in {"colab_link", "github_link"}:
                correct_link = relative_notebook_path
            elif key == "colab_enterprise_link":
                correct_link = urllib.parse.quote(
                    f"{RAW_URL_PREFIX}{relative_notebook_path}",
                    safe=":",
                )
            elif key == "workbench_link":
                correct_link = f"{RAW_URL_PREFIX}{relative_notebook_path}"
            elif key == "bigquery_studio_link":
                correct_link = f"{GITHUB_URL_PREFIX}{relative_notebook_path}"
            elif key in {
                "linkedin_link",
                "bluesky_link",
                "twitter_link",
                "reddit_link",
                "facebook_link",
            }:
                correct_link = encoded_url

            if correct_link.lower() not in line.lower():
                line = line.replace(line[start_index:end_index], correct_link)
                changes_made = True

        new_lines.append(line)

    return "\n".join(new_lines), changes_made


def update_notebook_links(notebook_paths: list[str]) -> None:
    """Checks and fixes specific types of links in the provided list of notebooks."""
    print("Checking notebook links...")
    links_updated_count = 0

    for notebook_path in notebook_paths:
        # False positive
        if "vector-search-2-intro" in notebook_path:
            continue
        try:
            with open(notebook_path, encoding="utf-8") as f:
                notebook = nbformat.read(f, as_version=4)

            relative_notebook_path = os.path.relpath(
                notebook_path, start=os.getcwd()
            ).lower()
            notebook_modified = False

            for cell in notebook.cells:
                if (
                    cell.cell_type == "markdown"
                    and "<table" in cell.source
                    and "colab" in cell.source
                ):
                    updated_source, changes_made = fix_markdown_links(
                        cell.source, relative_notebook_path
                    )
                    if changes_made:
                        cell.source = updated_source
                        notebook_modified = True

            if notebook_modified:
                links_updated_count += 1
                with open(notebook_path, "w", encoding="utf-8") as f:
                    nbformat.write(notebook, f)
                print(f"  -> Fixed links in {notebook_path}")

        except Exception as e:
            print(f"Could not check links in {notebook_path}. Error: {e}")

    if links_updated_count > 0:
        print(f"Fixed links in {links_updated_count} notebooks.")


def preprocess_notebook(
    notebook_paths: list[str],
    max_line_length: int = DEFAULT_RUFF_LINE_LENGTH,
) -> None:
    """Parses notebooks and wraps '@param' blocks with '# fmt: off/on'."""
    print("Checking for '@param' blocks to wrap with '# fmt: off/on'...")
    for path in notebook_paths:
        try:
            with open(path, encoding="utf-8") as f:
                notebook = nbformat.read(f, as_version=4)

            notebook_modified = False
            for cell in notebook.cells:
                if cell.cell_type != "code" or "@param" not in cell.source:
                    continue

                source_lines = cell.source.split("\n")

                clean_source_lines = [
                    line
                    for line in source_lines
                    if line.strip() not in ["# fmt: off", "# fmt: on"]
                ]

                param_indices = [
                    i
                    for i, line in enumerate(clean_source_lines)
                    if "@param" in line and len(line) >= max_line_length
                ]

                if not param_indices:
                    continue

                first_param_index = param_indices[0]
                last_param_index = param_indices[-1]

                # Reconstruct the cell source with new directives
                new_source_lines = []
                for i, line in enumerate(clean_source_lines):
                    if i == first_param_index:
                        new_source_lines.append("# fmt: off")
                    new_source_lines.append(line)
                    if i == last_param_index:
                        new_source_lines.append("# fmt: on")

                new_source = "\n".join(new_source_lines)

                if new_source != cell.source:
                    cell.source = new_source
                    notebook_modified = True

            if notebook_modified:
                print(f"  -> Wrapped '@param' blocks with '# fmt: off/on' in {path}")
                with open(path, "w", encoding="utf-8") as f:
                    nbformat.write(notebook, f)

        except Exception as e:
            print(f"Could not process notebook {path}. Error: {e}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(
            "Usage: python notebook_processor.py <notebook_path1> <notebook_path2> ..."
        )
        sys.exit(1)

    notebook_paths = sys.argv[1:]
    preprocess_notebook(notebook_paths)
    update_notebook_links(notebook_paths)