diff --git a/pyproject.toml b/pyproject.toml index 13fc5934f28d..cf5406acf199 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,9 @@ dependencies = [ "polyfile-weave>=0.5.9", ] +[project.scripts] +weave = "weave.cli.main:main" + [project.optional-dependencies] wandb = ["wandb>=0.17.1"] rich = ["rich"] # Optional dependency for enhanced console output diff --git a/tests/cli/test_login.py b/tests/cli/test_login.py new file mode 100644 index 000000000000..294a8eca5949 --- /dev/null +++ b/tests/cli/test_login.py @@ -0,0 +1,39 @@ +"""Tests for the Weave CLI login command.""" + +from unittest.mock import patch + +from click.testing import CliRunner + +from weave.cli.login import login as login_command + + +def test_cli_login_passes_normalized_host() -> None: + runner = CliRunner() + api_key = "a" * 40 + + with patch("weave.cli.login.wandb.login", return_value=True) as mock_login: + result = runner.invoke( + login_command, + [api_key, "--host", "api.wandb.ai", "--relogin"], + ) + + assert result.exit_code == 0 + mock_login.assert_called_once_with( + anonymous=None, + key=api_key, + relogin=True, + host="https://api.wandb.ai", + force=True, + timeout=None, + verify=False, + ) + + +def test_cli_login_failure_returns_error() -> None: + runner = CliRunner() + + with patch("weave.cli.login.wandb.login", return_value=False): + result = runner.invoke(login_command, []) + + assert result.exit_code != 0 + assert "Login failed." in result.output diff --git a/weave/cli/__init__.py b/weave/cli/__init__.py new file mode 100644 index 000000000000..f33a3bdf77bf --- /dev/null +++ b/weave/cli/__init__.py @@ -0,0 +1 @@ +"""Weave CLI package.""" diff --git a/weave/cli/login.py b/weave/cli/login.py new file mode 100644 index 000000000000..3855bf5ee71b --- /dev/null +++ b/weave/cli/login.py @@ -0,0 +1,78 @@ +"""Login command for the Weave CLI.""" + +from __future__ import annotations + +from typing import Literal + +import click + +from weave.compat import wandb +from weave.compat.wandb.wandb_thin.login import ( + _get_default_host as _compat_get_default_host, +) + + +def _get_default_host() -> str: + return _compat_get_default_host() + + +@click.command("login") +@click.argument("key", required=False) +@click.option( + "--anonymous", + type=click.Choice(["allow", "must", "never"], case_sensitive=False), + default=None, + help="Control anonymous login behavior.", +) +@click.option( + "--relogin", + is_flag=True, + default=False, + help="Force a relogin even if an API key is already configured.", +) +@click.option( + "--host", + default=None, + help="W&B host URL, e.g. https://api.wandb.ai", +) +@click.option( + "--verify/--no-verify", + default=False, + help="Verify the API key with the server.", +) +@click.option( + "--timeout", + type=int, + default=None, + help="Seconds to wait for user input.", +) +def login( + key: str | None, + anonymous: Literal["allow", "must", "never"] | None, + relogin: bool, + host: str | None, + verify: bool, + timeout: int | None, +) -> None: + """Log in to Weights & Biases for Weave.""" + normalized_host = host + if normalized_host is not None and not normalized_host.startswith( + ("http://", "https://") + ): + normalized_host = f"https://{normalized_host}" + + try: + success = wandb.login( + anonymous=anonymous, + key=key, + relogin=relogin, + host=normalized_host, + force=relogin, + timeout=timeout, + verify=verify, + ) + except Exception as exc: + raise click.ClickException(str(exc)) from exc + + if not success: + raise click.ClickException("Login failed.") diff --git a/weave/cli/main.py b/weave/cli/main.py new file mode 100644 index 000000000000..b90a9b1e6c82 --- /dev/null +++ b/weave/cli/main.py @@ -0,0 +1,24 @@ +"""Entry point for the Weave CLI.""" + +from __future__ import annotations + +import logging +from typing import cast + +import click + +from weave.cli.login import login as login_command + + +def _configure_logging() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + + +@click.group() +def main() -> None: + """Weave command line interface.""" + _configure_logging() + + +cli = cast(click.Group, main) +cli.add_command(login_command) diff --git a/weave/compat/wandb/wandb_thin/login.py b/weave/compat/wandb/wandb_thin/login.py index e7e4ca8e623e..ae4676405bb2 100644 --- a/weave/compat/wandb/wandb_thin/login.py +++ b/weave/compat/wandb/wandb_thin/login.py @@ -80,6 +80,16 @@ def _clear_setting(key: str) -> None: logger.warning(f"Failed to clear setting {key}: {e}") +def _normalize_host(host: str | None) -> str | None: + """Normalize a host or base URL to a hostname without scheme or trailing slash.""" + if host is None: + return None + host = host.rstrip("/") + if host.startswith(("http://", "https://")): + host = host.split("://", 1)[1] + return host + + def login( anonymous: Literal["must", "allow", "never"] | None = None, key: str | None = None, @@ -165,7 +175,8 @@ def __init__( self._force = force self._timeout = timeout self._key = key - self._host = host if host else _get_default_host() + resolved_host = host if host else _get_default_host() + self._host = _normalize_host(resolved_host) or resolved_host self.is_anonymous = anonymous == "must" def is_apikey_configured(self) -> bool: