google-colab-cli 0.5.4__py3-none-any.whl

colab_cli/auth.py

@@ -0,0 +1,170 @@
+# 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.
+
+import enum
+import json
+import logging
+import os
+import warnings
+from typing import Optional
+
+import google.auth
+from google.auth.transport import requests
+from google.auth.transport.requests import Request
+from google.oauth2.credentials import Credentials
+from google_auth_oauthlib.flow import InstalledAppFlow
+
+logger = logging.getLogger(__name__)
+
+
+class AuthProvider(str, enum.Enum):
+    """Authentication strategy for talking to the Colab backend.
+
+    Values are the lowercase strings accepted by the global ``--auth`` flag.
+    """
+
+    OAUTH2 = "oauth2"
+    ADC = "adc"
+
+
+# Standard Scopes for Colab and Drive (Public Auth)
+PUBLIC_SCOPES = [
+    "openid",
+    "https://www.googleapis.com/auth/userinfo.profile",
+    "https://www.googleapis.com/auth/userinfo.email",
+    "https://www.googleapis.com/auth/colaboratory",
+    "https://www.googleapis.com/auth/drive.file",
+]
+
+
+TOKEN_CONFIG_PATH = os.path.expanduser("~/.config/colab-cli/token.json")
+OAUTH_SERVER_PORT = 8200
+
+
+def _get_google_auth_credentials(config_path: str) -> Credentials:
+    """
+    Retrieves credentials using standard public OAuth2 flow.
+    """
+    client_config = None
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            client_config = json.load(f)
+    if not client_config:
+        raise FileNotFoundError(
+            f"Client OAuth config not found at {config_path}. "
+            "Please provide a valid path via -c/--client-oauth-config."
+        )
+
+    creds = None
+
+    # Ensure config directory exists for the token file
+    os.makedirs(os.path.dirname(TOKEN_CONFIG_PATH), exist_ok=True)
+
+    if os.path.exists(TOKEN_CONFIG_PATH):
+        try:
+            creds = Credentials.from_authorized_user_file(
+                TOKEN_CONFIG_PATH, PUBLIC_SCOPES
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load token from {TOKEN_CONFIG_PATH}: {e}")
+
+    if not creds or not creds.valid:
+        if creds and creds.expired and creds.refresh_token:
+            try:
+                creds.refresh(Request())
+            except Exception as e:
+                logger.warning(f"Failed to refresh token: {e}")
+                creds = None
+
+        if not creds:
+            flow = InstalledAppFlow.from_client_config(client_config, PUBLIC_SCOPES)
+            creds = flow.run_local_server(port=OAUTH_SERVER_PORT)
+
+        # Save the credentials for the next run
+        try:
+            with open(TOKEN_CONFIG_PATH, "w") as token_file:
+                token_file.write(creds.to_json())
+        except Exception as e:
+            logger.error(f"Failed to save token to {TOKEN_CONFIG_PATH}: {e}")
+
+    return creds
+
+
+def _get_adc_credentials() -> Credentials:
+    """Retrieves credentials using Google Application Default Credentials.
+
+    Honors the standard ADC discovery chain (``GOOGLE_APPLICATION_CREDENTIALS``,
+    ``gcloud auth application-default login``, GCE/GKE metadata server, etc.).
+
+    The RuntimeService at colab.pa.googleapis.com requires the
+    `colaboratory` scope (otherwise keep-alive returns 403 SCOPE_NOT_PERMITTED).
+    Most ADC credential types (service accounts, GCE/GKE, impersonated)
+    support `with_scopes`; user credentials minted by
+    `gcloud auth application-default login` do not. For the latter, the user
+    must re-run `gcloud auth application-default login` with
+    `--scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory`
+    (`openid` and `cloud-platform` are required by `gcloud` itself; `userinfo.email`
+    is required by the session backend; `colaboratory` is required by this RPC).
+    """
+    # `google.auth._default` emits a UserWarning when ADC user credentials
+    # don't have a quota project pinned ("Your application has authenticated
+    # using end user credentials from Google Cloud SDK without a quota
+    # project. You might receive a 'quota exceeded' or 'API not enabled'
+    # error.").
+    #
+    # That heuristic does not apply to this CLI: every call we make to
+    # `colab.pa.googleapis.com` carries `X-Goog-User-Project: 1014160490159`
+    # (Colab's project id) — see AGENTS.md item 18 — so the user's
+    # quota-project setting is irrelevant. The warning shows up on every
+    # single `colab` invocation under ADC, which is pure noise. Filter it,
+    # but keep the scope as tight as possible: only this exact message,
+    # only during this one call.
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"Your application has authenticated using end user credentials.*",
+            category=UserWarning,
+        )
+        creds, _ = google.auth.default(scopes=list(PUBLIC_SCOPES))
+    # Some credential subclasses ignore the `scopes=` kwarg in `default()`
+    # (e.g. user creds), so re-apply via `with_scopes` when supported.
+    if getattr(creds, "requires_scopes", False):
+        try:
+            creds = creds.with_scopes(list(PUBLIC_SCOPES))
+        except Exception as e:  # NotImplementedError for non-scopable creds.
+            logger.debug(f"Could not augment ADC scopes via with_scopes: {e}")
+    return creds
+
+
+def get_credentials(
+    config_path: Optional[str] = None,
+    provider: AuthProvider = AuthProvider.OAUTH2,
+) -> requests.AuthorizedSession:
+    """Unified entry point for retrieving an authorized session.
+
+    Args:
+      config_path: Path to the OAuth2 client config JSON. Only consulted when
+        ``provider`` is ``OAUTH2``.
+      provider: Which authentication strategy to use.
+    """
+    if provider == AuthProvider.OAUTH2:
+        if not config_path:
+            config_path = os.path.expanduser("~/.colab-cli-oauth-config.json")
+        creds = _get_google_auth_credentials(config_path)
+    elif provider == AuthProvider.ADC:
+        creds = _get_adc_credentials()
+    else:
+        raise ValueError(f"Unknown auth provider: {provider!r}")
+
+    return requests.AuthorizedSession(creds)

colab_cli/auto_update.py

@@ -0,0 +1,226 @@
+# 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.
+
+"""Auto-update subsystem.
+
+Owns version detection, the PyPI-style update probe, the on-disk
+``latest_version`` cache, and the upgrade-banner UX. The CLI's global
+callback (``cli.py``) calls ``check_for_updates`` once per day and
+``maybe_show_cached_banner`` on every other invocation; the
+``colab update`` Typer command (``commands/utility.py``) delegates to
+``check_for_updates``.
+"""
+
+import json
+import subprocess
+import urllib.request
+from datetime import datetime, timezone
+from importlib.metadata import PackageNotFoundError, version as installed_version
+from packaging.version import InvalidVersion, Version
+from typing import Optional
+
+import typer
+
+from colab_cli.common import state
+from colab_cli.state import Settings
+
+
+# ---------- Version detection -------------------------------------------
+
+
+def get_app_version() -> str:
+    """Return the installed package version, falling back to the git short hash."""
+    try:
+        return installed_version("colab")
+    except (PackageNotFoundError, InvalidVersion):
+        pass
+
+    try:
+        return subprocess.check_output(
+            ["git", "rev-parse", "--short", "HEAD"],
+            stderr=subprocess.DEVNULL,
+            encoding="utf-8",
+        ).strip()
+    except Exception:
+        return "unknown"
+
+
+# ---------- Source fetchers ---------------------------------------------
+
+
+def _parse_version(payload: Optional[dict]) -> Optional[str]:
+    """Returns ``info.version`` from a PyPI-style payload, or None."""
+    return (payload or {}).get("info", {}).get("version")
+
+
+def _fetch_pypi(url: str, quiet: bool) -> Optional[dict]:
+    """Fetches and parses the PyPI-style JSON document at ``url``."""
+    try:
+        with urllib.request.urlopen(url, timeout=5) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except Exception as e:
+        if not quiet:
+            typer.echo(f"[colab] Warning: Failed to fetch update info: {e}")
+        return None
+
+
+# ---------- Version comparison ------------------------------------------
+
+
+def _is_newer(candidate: Optional[str], current: str) -> bool:
+    """True when ``candidate`` strictly exceeds ``current`` (PEP 440)."""
+    if not candidate:
+        return False
+    try:
+        return Version(candidate) > Version(current)
+    except InvalidVersion:
+        return candidate != current
+
+
+# ---------- UX ----------------------------------------------------------
+
+
+def announce_upgrade(
+    latest: str,
+    current: str,
+    install_cmd: str,
+    *,
+    show_disable_hint: bool = False,
+) -> None:
+    """Print the upgrade banner.
+
+    ``show_disable_hint`` controls whether the trailing line that explains
+    how to silence the auto-check is included. It is only added when the
+    banner is shown unsolicited (the daily background fetch and the cached
+    banner on subsequent invocations); explicit ``colab update`` calls
+    omit it because the user already opted in to seeing the result.
+    """
+    typer.echo(
+        f"\n[colab] A new version of Colab CLI is available: {latest} (current: {current})"
+    )
+    typer.echo(f"[colab] Run '{install_cmd}' to update.")
+    if show_disable_hint:
+        typer.echo(
+            "[colab] To silence this check, set "
+            '"enable_update_check": false in '
+            "~/.config/colab-cli/settings.json"
+        )
+    typer.echo("")
+
+
+# ---------- Orchestration -----------------------------------------------
+
+
+def check_for_updates(quiet: bool = False) -> None:
+    """Check PyPI for updates and print a message if a new version is available.
+
+    The disable-hint is appended to the banner only when ``quiet`` is True
+    (the daily background fetch); explicit ``colab update`` invocations
+    (``quiet=False``) omit it because the user asked for the check.
+    """
+    settings = state.settings_store.load()
+    current = get_app_version()
+
+    try:
+        pypi = _fetch_pypi(settings.update_url, quiet)
+        pypi_v = _parse_version(pypi)
+
+        if _is_newer(pypi_v, current):
+            announce_upgrade(
+                pypi_v,
+                current,
+                "pip install --upgrade colab",
+                show_disable_hint=quiet,
+            )
+        elif not quiet:
+            suffix = f", latest: {pypi_v}" if pypi_v else ""
+            typer.echo(f"[colab] Colab CLI is up to date (version: {current}{suffix}).")
+
+        # Cache the highest observed version; never downgrade.
+        cached = settings.latest_version or "0"
+        if _is_newer(pypi_v, cached):
+            settings.latest_version = pypi_v
+
+        settings.last_check = datetime.now(timezone.utc)
+        state.settings_store.save(settings)
+
+    except Exception as e:
+        if not quiet:
+            typer.echo(f"[colab] Failed to check for updates: {e}")
+
+
+# ---------- Background hooks (called from cli.py) -----------------------
+
+
+def _is_throttled(settings: Settings, *, now: Optional[datetime] = None) -> bool:
+    """True if the once-per-day fetch should be skipped."""
+    if settings.last_check is None:
+        return False
+    now = now or datetime.now(timezone.utc)
+    return (now - settings.last_check).days < 1
+
+
+def maybe_show_cached_banner(settings: Settings) -> None:
+    """Print the cached upgrade banner if the cache reports a newer version.
+
+    Called from the global CLI callback when the daily fetch is throttled.
+    The banner uses a generic ``colab update`` install hint because the
+    cache does not record which source supplied the version; the disable
+    hint is shown because this is unsolicited output.
+    """
+    if not settings.latest_version:
+        return
+    current = get_app_version()
+    if not _is_newer(settings.latest_version, current):
+        return
+    announce_upgrade(
+        settings.latest_version,
+        current,
+        "colab update",
+        show_disable_hint=True,
+    )
+
+
+def run_background_check() -> None:
+    """Entry point for the global CLI callback.
+
+    Performs either the daily fetch (which writes the cache) or, if
+    throttled, surfaces the cached banner. Honors the
+    ``enable_update_check`` master switch.
+    """
+    settings = state.settings_store.load()
+    if not settings.enable_update_check:
+        return
+    if _is_throttled(settings):
+        maybe_show_cached_banner(settings)
+    else:
+        check_for_updates(quiet=True)
+
+
+# ---------- Self-install ------------------------------------------------
+
+
+# PyPI distribution name (different from the importable package name `colab`).
+PYPI_PACKAGE_NAME = "google-colab-cli"
+
+
+def self_install() -> None:
+    """Run ``pip install -U <PYPI_PACKAGE_NAME>`` to upgrade the CLI in place."""
+    import sys
+
+    cmd = [sys.executable, "-m", "pip", "install", "-U", PYPI_PACKAGE_NAME]
+    typer.echo(f"[colab] Running: {' '.join(cmd)}")
+    result = subprocess.run(cmd)
+    if result.returncode != 0:
+        raise typer.Exit(code=result.returncode)

colab_cli/cli.py

@@ -0,0 +1,151 @@
+# 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.
+
+import os
+from typing import Optional
+
+import click
+import typer
+from typer.core import TyperGroup
+from typing_extensions import Annotated
+
+from colab_cli import auto_update
+from colab_cli.auth import AuthProvider
+from colab_cli.common import state, setup_logging
+from colab_cli.commands import session, execution, files, automation, run, utility
+
+
+class AlphabeticalGroup(TyperGroup):
+    """A `TyperGroup` that lists subcommands alphabetically in `--help` output.
+
+    Subcommands are registered in functional groups (session, execution, files,
+    automation, utility), but users discovering the CLI via `colab --help` /
+    `colab help` benefit from a deterministic, alphabetical listing.
+    """
+
+    def list_commands(self, ctx: click.Context) -> list[str]:
+        return sorted(super().list_commands(ctx))
+
+
+app = typer.Typer(
+    help="Colab CLI",
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]},
+    cls=AlphabeticalGroup,
+)
+
+
+@app.callback()
+def callback(
+    ctx: typer.Context,
+    client_oauth_config: Annotated[
+        str,
+        typer.Option(
+            "-c", "--client-oauth-config", help="Path to client OAuth config JSON file"
+        ),
+    ] = os.path.expanduser("~/.colab-cli-oauth-config.json"),
+    config: Annotated[
+        Optional[str],
+        typer.Option(
+            "--config",
+            help="Path to session state file (~/.config/colab-cli/sessions.json)",
+        ),
+    ] = None,
+    logtostderr: Annotated[
+        bool, typer.Option("--logtostderr", help="Log all output to stderr")
+    ] = False,
+    auth: Annotated[
+        AuthProvider,
+        typer.Option(
+            "--auth",
+            help=(
+                "Authentication strategy to use: 'oauth2' (public InstalledAppFlow),"
+                " or 'adc' (Application Default Credentials)."
+            ),
+            case_sensitive=False,
+        ),
+    ] = AuthProvider.ADC,
+):
+    """
+    Colab CLI global configuration.
+    """
+    state.client_oauth_config = client_oauth_config
+    state.config_path = config
+    state.logtostderr = logtostderr
+    state.auth_provider = auth
+    setup_logging(logtostderr)
+
+    # Daily fetch + cached banner on every invocation.
+    #
+    # Suppress the banner for short-lived informational subcommands so their
+    # output stays clean and machine-parseable:
+    #   - `update`: runs its own check + announce; would duplicate the banner.
+    #   - `version`, `log`, `pay`, `help`, `url`: pure-display commands whose
+    #     output users routinely pipe / scrape (e.g. `colab url -s s1 | xclip`);
+    #     a stochastic upgrade banner injected once a day would corrupt those
+    #     pipelines.
+    #   - `whoami`: developer-only debugging tool; banner would obscure the
+    #     auth/scope info the user invoked it to see.
+    _AUTO_UPDATE_SUPPRESSED = {
+        "update",
+        "version",
+        "log",
+        "pay",
+        "help",
+        "url",
+        "whoami",
+    }
+    if ctx.invoked_subcommand not in _AUTO_UPDATE_SUPPRESSED:
+        auto_update.run_background_check()
+
+
+@app.command(name="help")
+def help_command(
+    ctx: typer.Context,
+    command: Annotated[
+        Optional[str], typer.Argument(help="Command to show help for")
+    ] = None,
+):
+    """
+    Show help for a command.
+    """
+    if not command:
+        typer.echo(ctx.parent.get_help())
+        return
+
+    group = ctx.parent.command
+    cmd = group.get_command(ctx, command)
+    if cmd is None:
+        typer.echo(f"No such command '{command}'.", err=True)
+        raise typer.Exit(code=2)
+
+    with click.Context(cmd, info_name=command, parent=ctx.parent) as cmd_ctx:
+        typer.echo(cmd.get_help(cmd_ctx))
+
+
+# Register subcommands
+session.register(app)
+execution.register(app)
+files.register(app)
+automation.register(app)
+run.register(app)
+utility.register(app)
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()