chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,207 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user