From 19fb333b60fd7a739150f71fc29c68b46a4f132d Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Mon, 4 May 2026 13:38:48 -0400 Subject: [PATCH 1/4] feat(cli): show version update notice Signed-off-by: Eric W. Tramel --- packages/data-designer/pyproject.toml | 1 + .../src/data_designer/cli/main.py | 10 +- .../data-designer/src/data_designer/cli/ui.py | 23 ++ .../src/data_designer/cli/version_notice.py | 209 ++++++++++++++++++ packages/data-designer/tests/cli/test_main.py | 21 +- .../tests/cli/test_version_notice.py | 151 +++++++++++++ uv.lock | 4 +- 7 files changed, 416 insertions(+), 3 deletions(-) create mode 100644 packages/data-designer/src/data_designer/cli/version_notice.py create mode 100644 packages/data-designer/tests/cli/test_version_notice.py diff --git a/packages/data-designer/pyproject.toml b/packages/data-designer/pyproject.toml index 04ad73d97..df6cb2e2d 100644 --- a/packages/data-designer/pyproject.toml +++ b/packages/data-designer/pyproject.toml @@ -41,6 +41,7 @@ path = "dev-tools/hatch_build.py" dependencies = [ "data-designer-config=={{ version }}", "data-designer-engine=={{ version }}", + "packaging>=25,<27", "prompt-toolkit>=3.0.0,<4", "typer>=0.12.0,<1", ] diff --git a/packages/data-designer/src/data_designer/cli/main.py b/packages/data-designer/src/data_designer/cli/main.py index a6e68f3fa..c85a73cae 100644 --- a/packages/data-designer/src/data_designer/cli/main.py +++ b/packages/data-designer/src/data_designer/cli/main.py @@ -20,10 +20,18 @@ def _version_callback(value: bool) -> None: if not value: return try: - typer.echo(importlib.metadata.version(_PACKAGE_NAME)) + installed_version = importlib.metadata.version(_PACKAGE_NAME) except importlib.metadata.PackageNotFoundError: typer.echo(f"Unable to resolve installed {_PACKAGE_NAME} package version.", err=True) raise typer.Exit(1) from None + + typer.echo(installed_version) + from data_designer.cli.ui import print_update_notice + from data_designer.cli.version_notice import get_update_notice + + notice = get_update_notice(installed_version) + if notice is not None: + print_update_notice(notice.latest_version, notice.upgrade_command) raise typer.Exit() diff --git a/packages/data-designer/src/data_designer/cli/ui.py b/packages/data-designer/src/data_designer/cli/ui.py index eced52a7a..49b455df2 100644 --- a/packages/data-designer/src/data_designer/cli/ui.py +++ b/packages/data-designer/src/data_designer/cli/ui.py @@ -19,6 +19,7 @@ from rich.console import Console from rich.padding import Padding from rich.panel import Panel +from rich.text import Text from data_designer.config.utils.constants import RICH_CONSOLE_THEME, NordColor @@ -574,6 +575,28 @@ def print_info(message: str) -> None: _print_with_padding(f"💡 {message}") +def print_update_notice(latest_version: str, upgrade_command: str) -> None: + """Print a compact version update notice. + + Args: + latest_version: Latest available Data Designer version. + upgrade_command: Command users can run to upgrade. + """ + content = Text.assemble( + "New Data Designer version available: ", + (latest_version, f"bold {NordColor.NORD14.value}"), + "\nUpgrade with: ", + (upgrade_command, f"bold {NordColor.NORD8.value}"), + ) + panel = Panel.fit( + content, + title="🚀 Update available", + title_align="left", + border_style=NordColor.NORD8.value, + ) + _console.print(panel) + + def print_text(message: str) -> None: """Print a text message. diff --git a/packages/data-designer/src/data_designer/cli/version_notice.py b/packages/data-designer/src/data_designer/cli/version_notice.py new file mode 100644 index 000000000..5bb941d0c --- /dev/null +++ b/packages/data-designer/src/data_designer/cli/version_notice.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import sys +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from packaging.version import InvalidVersion, Version + +from data_designer.config.utils.constants import DATA_DESIGNER_HOME + +_PACKAGE_NAME = "data-designer" +_PYPI_JSON_URL = f"https://pypi.org/pypi/{_PACKAGE_NAME}/json" +_VERSION_CHECK_TIMEOUT_SECONDS = 0.75 +_CACHE_TTL_SECONDS = 6 * 60 * 60 +_CACHE_FILE_NAME = "version-check.json" +_DISABLE_VERSION_CHECK_ENV_VAR = "DATA_DESIGNER_DISABLE_VERSION_CHECK" +_INCLUDE_PRERELEASES_ENV_VAR = "DATA_DESIGNER_VERSION_CHECK_PRERELEASES" +_UV_TOOL_UPGRADE_COMMAND = "uv tool upgrade data-designer" +_PROJECT_UPGRADE_COMMAND = "uv add --upgrade data-designer" + + +@dataclass(frozen=True) +class UpdateNotice: + latest_version: str + upgrade_command: str + + +def get_update_notice( + installed_version: str, + *, + cache_dir: Path = DATA_DESIGNER_HOME, + environ: Mapping[str, str] | None = None, + now: Callable[[], float] = time.time, + python_prefix: str | None = None, +) -> UpdateNotice | None: + env = os.environ if environ is None else environ + if _env_flag_enabled(env, _DISABLE_VERSION_CHECK_ENV_VAR): + return None + + try: + installed = Version(installed_version) + except InvalidVersion: + return None + + include_prereleases = installed.is_prerelease or _env_flag_enabled(env, _INCLUDE_PRERELEASES_ENV_VAR) + latest_version = _get_latest_version( + include_prereleases=include_prereleases, + cache_dir=cache_dir, + now=now, + ) + if latest_version is None: + return None + + try: + latest = Version(latest_version) + except InvalidVersion: + return None + + if latest <= installed: + return None + + return UpdateNotice( + latest_version=latest.public, + upgrade_command=select_upgrade_command(environ=env, python_prefix=python_prefix), + ) + + +def select_upgrade_command( + *, + environ: Mapping[str, str] | None = None, + python_prefix: str | None = None, +) -> str: + env = os.environ if environ is None else environ + prefix = Path(sys.prefix if python_prefix is None else python_prefix) + prefix_parts = set(prefix.parts) + if "uv" in prefix_parts and "tools" in prefix_parts: + return _UV_TOOL_UPGRADE_COMMAND + if env.get("UV_PROJECT_ENVIRONMENT") or env.get("VIRTUAL_ENV"): + return _PROJECT_UPGRADE_COMMAND + if prefix.name == ".venv": + return _PROJECT_UPGRADE_COMMAND + return _UV_TOOL_UPGRADE_COMMAND + + +def _get_latest_version( + *, + include_prereleases: bool, + cache_dir: Path, + now: Callable[[], float], +) -> str | None: + cache_path = cache_dir / _CACHE_FILE_NAME + cached_version = _read_cached_version( + cache_path=cache_path, + include_prereleases=include_prereleases, + now=now, + ) + if cached_version is not None: + return cached_version + + try: + latest_version = _fetch_latest_version(include_prereleases=include_prereleases) + except (HTTPError, URLError, TimeoutError, OSError, json.JSONDecodeError): + return None + + if latest_version is not None: + _write_cached_version( + cache_path=cache_path, + latest_version=latest_version, + include_prereleases=include_prereleases, + checked_at=now(), + ) + return latest_version + + +def _fetch_latest_version(*, include_prereleases: bool) -> str | None: + request = Request(_PYPI_JSON_URL, headers={"Accept": "application/json", "User-Agent": "data-designer"}) + with urlopen(request, timeout=_VERSION_CHECK_TIMEOUT_SECONDS) as response: + payload = json.load(response) + if not isinstance(payload, dict): + return None + return _latest_version_from_pypi_payload(payload, include_prereleases=include_prereleases) + + +def _latest_version_from_pypi_payload(payload: Mapping[str, Any], *, include_prereleases: bool) -> str | None: + releases = payload.get("releases") + if not isinstance(releases, dict): + return None + + candidates: list[Version] = [] + for version_text, release_files in releases.items(): + if not isinstance(version_text, str) or _is_yanked_release(release_files): + continue + try: + version = Version(version_text) + except InvalidVersion: + continue + if version.is_prerelease and not include_prereleases: + continue + candidates.append(version) + + if not candidates: + return None + + return max(candidates).public + + +def _is_yanked_release(release_files: Any) -> bool: + if not isinstance(release_files, list) or not release_files: + return True + return all(isinstance(release_file, dict) and release_file.get("yanked", False) for release_file in release_files) + + +def _read_cached_version( + *, + cache_path: Path, + include_prereleases: bool, + now: Callable[[], float], +) -> str | None: + try: + cache_data = json.loads(cache_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + if not isinstance(cache_data, dict): + return None + + if cache_data.get("include_prereleases") != include_prereleases: + return None + + checked_at = cache_data.get("checked_at") + latest_version = cache_data.get("latest_version") + if not isinstance(checked_at, (int, float)) or not isinstance(latest_version, str): + return None + if now() - float(checked_at) > _CACHE_TTL_SECONDS: + return None + return latest_version + + +def _write_cached_version( + *, + cache_path: Path, + latest_version: str, + include_prereleases: bool, + checked_at: float, +) -> None: + cache_data = { + "checked_at": checked_at, + "include_prereleases": include_prereleases, + "latest_version": latest_version, + } + try: + cache_path.parent.mkdir(parents=True, exist_ok=True) + cache_path.write_text(json.dumps(cache_data), encoding="utf-8") + except OSError: + return + + +def _env_flag_enabled(env: Mapping[str, str], name: str) -> bool: + value = env.get(name, "") + return value.strip().lower() in {"1", "true", "yes", "on"} diff --git a/packages/data-designer/tests/cli/test_main.py b/packages/data-designer/tests/cli/test_main.py index 32d9cfc7d..eac451bce 100644 --- a/packages/data-designer/tests/cli/test_main.py +++ b/packages/data-designer/tests/cli/test_main.py @@ -9,6 +9,7 @@ from typer.testing import CliRunner from data_designer.cli.main import app, main +from data_designer.cli.version_notice import UpdateNotice from data_designer.config.utils.constants import DEFAULT_NUM_RECORDS runner = CliRunner() @@ -51,12 +52,30 @@ def test_main_skips_bootstrap_when_version_follows_another_flag(mock_bootstrap: def test_app_version_prints_installed_data_designer_version() -> None: - with patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0") as mock_version: + with ( + patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0") as mock_version, + patch("data_designer.cli.version_notice.get_update_notice", return_value=None) as mock_notice, + ): result = runner.invoke(app, ["--version"]) assert result.exit_code == 0 assert result.output == "0.6.0\n" mock_version.assert_called_once_with("data-designer") + mock_notice.assert_called_once_with("0.6.0") + + +def test_app_version_prints_update_notice_after_installed_version() -> None: + notice = UpdateNotice(latest_version="0.6.1", upgrade_command="uv tool upgrade data-designer") + with ( + patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0"), + patch("data_designer.cli.version_notice.get_update_notice", return_value=notice), + ): + result = runner.invoke(app, ["--version"]) + + assert result.exit_code == 0 + assert result.output.splitlines()[0] == "0.6.0" + assert "New Data Designer version available: 0.6.1" in result.output + assert "Upgrade with: uv tool upgrade data-designer" in result.output def test_app_version_errors_when_package_version_is_missing() -> None: diff --git a/packages/data-designer/tests/cli/test_version_notice.py b/packages/data-designer/tests/cli/test_version_notice.py new file mode 100644 index 000000000..119bfdf1d --- /dev/null +++ b/packages/data-designer/tests/cli/test_version_notice.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import Mock + +from pytest import MonkeyPatch + +from data_designer.cli import version_notice +from data_designer.cli.version_notice import get_update_notice, select_upgrade_command + + +def test_get_update_notice_returns_notice_for_newer_stable_version(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(return_value="0.6.1") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice( + "0.6.0", + cache_dir=tmp_path, + environ={}, + now=lambda: 1_000.0, + python_prefix="/opt/python", + ) + + assert notice is not None + assert notice.latest_version == "0.6.1" + assert notice.upgrade_command == "uv tool upgrade data-designer" + mock_fetch.assert_called_once_with(include_prereleases=False) + + +def test_get_update_notice_returns_none_for_current_version(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(return_value="0.6.0") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.0", cache_dir=tmp_path, environ={}, now=lambda: 1_000.0) + + assert notice is None + mock_fetch.assert_called_once_with(include_prereleases=False) + + +def test_get_update_notice_fails_closed_when_check_fails(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(side_effect=OSError("network unavailable")) + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.0", cache_dir=tmp_path, environ={}, now=lambda: 1_000.0) + + assert notice is None + mock_fetch.assert_called_once_with(include_prereleases=False) + + +def test_get_update_notice_respects_opt_out(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(return_value="0.6.1") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice( + "0.6.0", + cache_dir=tmp_path, + environ={"DATA_DESIGNER_DISABLE_VERSION_CHECK": "1"}, + now=lambda: 1_000.0, + ) + + assert notice is None + mock_fetch.assert_not_called() + + +def test_get_update_notice_uses_fresh_cache(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + cache_path = tmp_path / "version-check.json" + cache_path.write_text( + json.dumps( + { + "checked_at": 1_000.0, + "include_prereleases": False, + "latest_version": "0.6.1", + } + ), + encoding="utf-8", + ) + mock_fetch = Mock(return_value="0.6.2") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.0", cache_dir=tmp_path, environ={}, now=lambda: 1_001.0) + + assert notice is not None + assert notice.latest_version == "0.6.1" + mock_fetch.assert_not_called() + + +def test_prerelease_versions_are_ignored_unless_requested() -> None: + payload = { + "releases": { + "0.6.1": [{"yanked": False}], + "0.6.2rc1": [{"yanked": False}], + } + } + + assert version_notice._latest_version_from_pypi_payload(payload, include_prereleases=False) == "0.6.1" + assert version_notice._latest_version_from_pypi_payload(payload, include_prereleases=True) == "0.6.2rc1" + + +def test_installed_prerelease_opts_into_prerelease_checks(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(return_value="0.6.2rc2") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.2rc1", cache_dir=tmp_path, environ={}, now=lambda: 1_000.0) + + assert notice is not None + assert notice.latest_version == "0.6.2rc2" + mock_fetch.assert_called_once_with(include_prereleases=True) + + +def test_prerelease_environment_flag_opts_into_prerelease_checks(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + mock_fetch = Mock(return_value="0.6.2rc1") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice( + "0.6.1", + cache_dir=tmp_path, + environ={"DATA_DESIGNER_VERSION_CHECK_PRERELEASES": "true"}, + now=lambda: 1_000.0, + ) + + assert notice is not None + assert notice.latest_version == "0.6.2rc1" + mock_fetch.assert_called_once_with(include_prereleases=True) + + +def test_select_upgrade_command_defaults_to_uv_tool() -> None: + command = select_upgrade_command(environ={}, python_prefix="/opt/python") + + assert command == "uv tool upgrade data-designer" + + +def test_select_upgrade_command_detects_uv_tool_environment() -> None: + command = select_upgrade_command( + environ={"VIRTUAL_ENV": "/repo/.venv"}, + python_prefix="/Users/user/.local/share/uv/tools/data-designer", + ) + + assert command == "uv tool upgrade data-designer" + + +def test_select_upgrade_command_detects_project_environment() -> None: + command = select_upgrade_command( + environ={"UV_PROJECT_ENVIRONMENT": ".venv"}, + python_prefix="/repo/.venv", + ) + + assert command == "uv add --upgrade data-designer" diff --git a/uv.lock b/uv.lock index 37e78f781..1ba121034 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", @@ -801,6 +801,7 @@ source = { editable = "packages/data-designer" } dependencies = [ { name = "data-designer-config" }, { name = "data-designer-engine" }, + { name = "packaging" }, { name = "prompt-toolkit" }, { name = "typer" }, ] @@ -809,6 +810,7 @@ dependencies = [ requires-dist = [ { name = "data-designer-config", editable = "packages/data-designer-config" }, { name = "data-designer-engine", editable = "packages/data-designer-engine" }, + { name = "packaging", specifier = ">=25,<27" }, { name = "prompt-toolkit", specifier = ">=3.0.0,<4" }, { name = "typer", specifier = ">=0.12.0,<1" }, ] From 959659339aa2e6dc9f03eccb00fea3a5d6dad633 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Mon, 4 May 2026 13:59:06 -0400 Subject: [PATCH 2/4] fix(cli): fail closed on version update lookup Signed-off-by: Eric W. Tramel --- packages/data-designer/src/data_designer/cli/main.py | 6 +++++- packages/data-designer/tests/cli/test_main.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/packages/data-designer/src/data_designer/cli/main.py b/packages/data-designer/src/data_designer/cli/main.py index c85a73cae..2753d5c9c 100644 --- a/packages/data-designer/src/data_designer/cli/main.py +++ b/packages/data-designer/src/data_designer/cli/main.py @@ -29,7 +29,11 @@ def _version_callback(value: bool) -> None: from data_designer.cli.ui import print_update_notice from data_designer.cli.version_notice import get_update_notice - notice = get_update_notice(installed_version) + try: + # The update CTA is opportunistic; version output should stay usable if lookup fails. + notice = get_update_notice(installed_version) + except Exception: + notice = None if notice is not None: print_update_notice(notice.latest_version, notice.upgrade_command) raise typer.Exit() diff --git a/packages/data-designer/tests/cli/test_main.py b/packages/data-designer/tests/cli/test_main.py index eac451bce..48ca8af71 100644 --- a/packages/data-designer/tests/cli/test_main.py +++ b/packages/data-designer/tests/cli/test_main.py @@ -78,6 +78,17 @@ def test_app_version_prints_update_notice_after_installed_version() -> None: assert "Upgrade with: uv tool upgrade data-designer" in result.output +def test_app_version_skips_update_notice_when_lookup_fails() -> None: + with ( + patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0"), + patch("data_designer.cli.version_notice.get_update_notice", side_effect=RuntimeError("network failure")), + ): + result = runner.invoke(app, ["--version"]) + + assert result.exit_code == 0 + assert result.output == "0.6.0\n" + + def test_app_version_errors_when_package_version_is_missing() -> None: with patch( "data_designer.cli.main.importlib.metadata.version", From a5e7cd10bd57a77a20db947d4a5f99e91fac98e1 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Mon, 4 May 2026 14:06:30 -0400 Subject: [PATCH 3/4] fix(cli): harden version update notice Signed-off-by: Eric W. Tramel --- .../src/data_designer/cli/main.py | 9 ++ .../src/data_designer/cli/version_notice.py | 32 +++++- packages/data-designer/tests/cli/test_main.py | 16 +++ .../tests/cli/test_version_notice.py | 105 ++++++++++++++++++ 4 files changed, 156 insertions(+), 6 deletions(-) diff --git a/packages/data-designer/src/data_designer/cli/main.py b/packages/data-designer/src/data_designer/cli/main.py index 2753d5c9c..aae743c75 100644 --- a/packages/data-designer/src/data_designer/cli/main.py +++ b/packages/data-designer/src/data_designer/cli/main.py @@ -5,6 +5,7 @@ import importlib.metadata import sys +from typing import TextIO import typer @@ -16,6 +17,11 @@ _PACKAGE_NAME = "data-designer" +def _should_show_update_notice(stream: TextIO | None = None) -> bool: + stream = sys.stdout if stream is None else stream + return stream.isatty() + + def _version_callback(value: bool) -> None: if not value: return @@ -26,6 +32,9 @@ def _version_callback(value: bool) -> None: raise typer.Exit(1) from None typer.echo(installed_version) + if not _should_show_update_notice(): + raise typer.Exit() + from data_designer.cli.ui import print_update_notice from data_designer.cli.version_notice import get_update_notice diff --git a/packages/data-designer/src/data_designer/cli/version_notice.py b/packages/data-designer/src/data_designer/cli/version_notice.py index 5bb941d0c..e5ca20d5c 100644 --- a/packages/data-designer/src/data_designer/cli/version_notice.py +++ b/packages/data-designer/src/data_designer/cli/version_notice.py @@ -22,11 +22,13 @@ _PYPI_JSON_URL = f"https://pypi.org/pypi/{_PACKAGE_NAME}/json" _VERSION_CHECK_TIMEOUT_SECONDS = 0.75 _CACHE_TTL_SECONDS = 6 * 60 * 60 +_CACHE_SCHEMA_VERSION = 1 _CACHE_FILE_NAME = "version-check.json" _DISABLE_VERSION_CHECK_ENV_VAR = "DATA_DESIGNER_DISABLE_VERSION_CHECK" _INCLUDE_PRERELEASES_ENV_VAR = "DATA_DESIGNER_VERSION_CHECK_PRERELEASES" _UV_TOOL_UPGRADE_COMMAND = "uv tool upgrade data-designer" _PROJECT_UPGRADE_COMMAND = "uv add --upgrade data-designer" +_PIPX_UPGRADE_COMMAND = "pipx upgrade data-designer" @dataclass(frozen=True) @@ -51,6 +53,8 @@ def get_update_notice( installed = Version(installed_version) except InvalidVersion: return None + if installed.local is not None: + return None include_prereleases = installed.is_prerelease or _env_flag_enabled(env, _INCLUDE_PRERELEASES_ENV_VAR) latest_version = _get_latest_version( @@ -83,6 +87,8 @@ def select_upgrade_command( env = os.environ if environ is None else environ prefix = Path(sys.prefix if python_prefix is None else python_prefix) prefix_parts = set(prefix.parts) + if "pipx" in prefix_parts and "venvs" in prefix_parts: + return _PIPX_UPGRADE_COMMAND if "uv" in prefix_parts and "tools" in prefix_parts: return _UV_TOOL_UPGRADE_COMMAND if env.get("UV_PROJECT_ENVIRONMENT") or env.get("VIRTUAL_ENV"): @@ -138,7 +144,7 @@ def _latest_version_from_pypi_payload(payload: Mapping[str, Any], *, include_pre candidates: list[Version] = [] for version_text, release_files in releases.items(): - if not isinstance(version_text, str) or _is_yanked_release(release_files): + if not isinstance(version_text, str) or not _has_installable_release_file(release_files): continue try: version = Version(version_text) @@ -154,10 +160,12 @@ def _latest_version_from_pypi_payload(payload: Mapping[str, Any], *, include_pre return max(candidates).public -def _is_yanked_release(release_files: Any) -> bool: - if not isinstance(release_files, list) or not release_files: - return True - return all(isinstance(release_file, dict) and release_file.get("yanked", False) for release_file in release_files) +def _has_installable_release_file(release_files: Any) -> bool: + if not isinstance(release_files, list): + return False + return any( + isinstance(release_file, dict) and not release_file.get("yanked", False) for release_file in release_files + ) def _read_cached_version( @@ -173,6 +181,10 @@ def _read_cached_version( if not isinstance(cache_data, dict): return None + if cache_data.get("schema_version") != _CACHE_SCHEMA_VERSION: + return None + if cache_data.get("package_name") != _PACKAGE_NAME: + return None if cache_data.get("include_prereleases") != include_prereleases: return None @@ -196,11 +208,19 @@ def _write_cached_version( "checked_at": checked_at, "include_prereleases": include_prereleases, "latest_version": latest_version, + "package_name": _PACKAGE_NAME, + "schema_version": _CACHE_SCHEMA_VERSION, } + temp_path = cache_path.with_name(f"{cache_path.name}.{os.getpid()}.tmp") try: cache_path.parent.mkdir(parents=True, exist_ok=True) - cache_path.write_text(json.dumps(cache_data), encoding="utf-8") + temp_path.write_text(json.dumps(cache_data), encoding="utf-8") + temp_path.replace(cache_path) except OSError: + try: + temp_path.unlink() + except OSError: + pass return diff --git a/packages/data-designer/tests/cli/test_main.py b/packages/data-designer/tests/cli/test_main.py index 48ca8af71..e30624a89 100644 --- a/packages/data-designer/tests/cli/test_main.py +++ b/packages/data-designer/tests/cli/test_main.py @@ -54,6 +54,7 @@ def test_main_skips_bootstrap_when_version_follows_another_flag(mock_bootstrap: def test_app_version_prints_installed_data_designer_version() -> None: with ( patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0") as mock_version, + patch("data_designer.cli.main._should_show_update_notice", return_value=True), patch("data_designer.cli.version_notice.get_update_notice", return_value=None) as mock_notice, ): result = runner.invoke(app, ["--version"]) @@ -68,6 +69,7 @@ def test_app_version_prints_update_notice_after_installed_version() -> None: notice = UpdateNotice(latest_version="0.6.1", upgrade_command="uv tool upgrade data-designer") with ( patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0"), + patch("data_designer.cli.main._should_show_update_notice", return_value=True), patch("data_designer.cli.version_notice.get_update_notice", return_value=notice), ): result = runner.invoke(app, ["--version"]) @@ -78,9 +80,23 @@ def test_app_version_prints_update_notice_after_installed_version() -> None: assert "Upgrade with: uv tool upgrade data-designer" in result.output +def test_app_version_skips_update_notice_when_stdout_is_not_tty() -> None: + with ( + patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0"), + patch("data_designer.cli.main._should_show_update_notice", return_value=False), + patch("data_designer.cli.version_notice.get_update_notice") as mock_notice, + ): + result = runner.invoke(app, ["--version"]) + + assert result.exit_code == 0 + assert result.output == "0.6.0\n" + mock_notice.assert_not_called() + + def test_app_version_skips_update_notice_when_lookup_fails() -> None: with ( patch("data_designer.cli.main.importlib.metadata.version", return_value="0.6.0"), + patch("data_designer.cli.main._should_show_update_notice", return_value=True), patch("data_designer.cli.version_notice.get_update_notice", side_effect=RuntimeError("network failure")), ): result = runner.invoke(app, ["--version"]) diff --git a/packages/data-designer/tests/cli/test_version_notice.py b/packages/data-designer/tests/cli/test_version_notice.py index 119bfdf1d..09be660fd 100644 --- a/packages/data-designer/tests/cli/test_version_notice.py +++ b/packages/data-designer/tests/cli/test_version_notice.py @@ -51,6 +51,32 @@ def test_get_update_notice_fails_closed_when_check_fails(tmp_path: Path, monkeyp mock_fetch.assert_called_once_with(include_prereleases=False) +def test_get_update_notice_returns_none_for_invalid_installed_version( + tmp_path: Path, + monkeypatch: MonkeyPatch, +) -> None: + mock_fetch = Mock(return_value="0.6.1") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("not-a-version", cache_dir=tmp_path, environ={}, now=lambda: 1_000.0) + + assert notice is None + mock_fetch.assert_not_called() + + +def test_get_update_notice_returns_none_for_local_installed_version( + tmp_path: Path, + monkeypatch: MonkeyPatch, +) -> None: + mock_fetch = Mock(return_value="0.6.1") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.1.dev0+gabc1234", cache_dir=tmp_path, environ={}, now=lambda: 1_000.0) + + assert notice is None + mock_fetch.assert_not_called() + + def test_get_update_notice_respects_opt_out(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: mock_fetch = Mock(return_value="0.6.1") monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) @@ -74,6 +100,8 @@ def test_get_update_notice_uses_fresh_cache(tmp_path: Path, monkeypatch: MonkeyP "checked_at": 1_000.0, "include_prereleases": False, "latest_version": "0.6.1", + "package_name": "data-designer", + "schema_version": 1, } ), encoding="utf-8", @@ -88,6 +116,57 @@ def test_get_update_notice_uses_fresh_cache(tmp_path: Path, monkeypatch: MonkeyP mock_fetch.assert_not_called() +def test_get_update_notice_refetches_expired_cache(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + cache_path = tmp_path / "version-check.json" + cache_path.write_text( + json.dumps( + { + "checked_at": 1_000.0, + "include_prereleases": False, + "latest_version": "0.6.1", + "package_name": "data-designer", + "schema_version": 1, + } + ), + encoding="utf-8", + ) + mock_fetch = Mock(return_value="0.6.2") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice( + "0.6.0", + cache_dir=tmp_path, + environ={}, + now=lambda: 1_000.0 + (7 * 60 * 60), + ) + + assert notice is not None + assert notice.latest_version == "0.6.2" + mock_fetch.assert_called_once_with(include_prereleases=False) + + +def test_get_update_notice_ignores_cache_with_old_schema(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + cache_path = tmp_path / "version-check.json" + cache_path.write_text( + json.dumps( + { + "checked_at": 1_000.0, + "include_prereleases": False, + "latest_version": "0.6.1", + } + ), + encoding="utf-8", + ) + mock_fetch = Mock(return_value="0.6.2") + monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) + + notice = get_update_notice("0.6.0", cache_dir=tmp_path, environ={}, now=lambda: 1_001.0) + + assert notice is not None + assert notice.latest_version == "0.6.2" + mock_fetch.assert_called_once_with(include_prereleases=False) + + def test_prerelease_versions_are_ignored_unless_requested() -> None: payload = { "releases": { @@ -100,6 +179,23 @@ def test_prerelease_versions_are_ignored_unless_requested() -> None: assert version_notice._latest_version_from_pypi_payload(payload, include_prereleases=True) == "0.6.2rc1" +def test_latest_version_ignores_yanked_and_malformed_release_files() -> None: + payload = { + "releases": { + "0.6.1": [{"yanked": False}], + "0.6.2": [{"yanked": True}, "not-a-file-record"], + "0.6.3": [], + } + } + + assert version_notice._latest_version_from_pypi_payload(payload, include_prereleases=False) == "0.6.1" + + +def test_latest_version_returns_none_for_malformed_pypi_payload() -> None: + assert version_notice._latest_version_from_pypi_payload({}, include_prereleases=False) is None + assert version_notice._latest_version_from_pypi_payload({"releases": []}, include_prereleases=False) is None + + def test_installed_prerelease_opts_into_prerelease_checks(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: mock_fetch = Mock(return_value="0.6.2rc2") monkeypatch.setattr(version_notice, "_fetch_latest_version", mock_fetch) @@ -142,6 +238,15 @@ def test_select_upgrade_command_detects_uv_tool_environment() -> None: assert command == "uv tool upgrade data-designer" +def test_select_upgrade_command_detects_pipx_environment() -> None: + command = select_upgrade_command( + environ={}, + python_prefix="/Users/user/.local/pipx/venvs/data-designer", + ) + + assert command == "pipx upgrade data-designer" + + def test_select_upgrade_command_detects_project_environment() -> None: command = select_upgrade_command( environ={"UV_PROJECT_ENVIRONMENT": ".venv"}, From d4fb67f36bb5563041f76a1befbea059709188f9 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Mon, 4 May 2026 14:13:09 -0400 Subject: [PATCH 4/4] fix(cli): tighten upgrade command detection Signed-off-by: Eric W. Tramel --- .../src/data_designer/cli/version_notice.py | 19 +++++++++++++------ .../tests/cli/test_version_notice.py | 9 +++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/packages/data-designer/src/data_designer/cli/version_notice.py b/packages/data-designer/src/data_designer/cli/version_notice.py index e5ca20d5c..265f49c76 100644 --- a/packages/data-designer/src/data_designer/cli/version_notice.py +++ b/packages/data-designer/src/data_designer/cli/version_notice.py @@ -86,18 +86,25 @@ def select_upgrade_command( ) -> str: env = os.environ if environ is None else environ prefix = Path(sys.prefix if python_prefix is None else python_prefix) - prefix_parts = set(prefix.parts) - if "pipx" in prefix_parts and "venvs" in prefix_parts: + prefix_parts = prefix.parts + if env.get("UV_PROJECT_ENVIRONMENT") or prefix.name == ".venv": + return _PROJECT_UPGRADE_COMMAND + if _has_direct_child_path(prefix_parts, "pipx", "venvs"): return _PIPX_UPGRADE_COMMAND - if "uv" in prefix_parts and "tools" in prefix_parts: + if _has_direct_child_path(prefix_parts, "uv", "tools"): return _UV_TOOL_UPGRADE_COMMAND - if env.get("UV_PROJECT_ENVIRONMENT") or env.get("VIRTUAL_ENV"): - return _PROJECT_UPGRADE_COMMAND - if prefix.name == ".venv": + if env.get("VIRTUAL_ENV"): return _PROJECT_UPGRADE_COMMAND return _UV_TOOL_UPGRADE_COMMAND +def _has_direct_child_path(parts: tuple[str, ...], parent: str, child: str) -> bool: + return any( + parts[index] == parent and parts[index + 1] == child and index + 2 == len(parts) - 1 + for index in range(len(parts) - 1) + ) + + def _get_latest_version( *, include_prereleases: bool, diff --git a/packages/data-designer/tests/cli/test_version_notice.py b/packages/data-designer/tests/cli/test_version_notice.py index 09be660fd..27add52c7 100644 --- a/packages/data-designer/tests/cli/test_version_notice.py +++ b/packages/data-designer/tests/cli/test_version_notice.py @@ -238,6 +238,15 @@ def test_select_upgrade_command_detects_uv_tool_environment() -> None: assert command == "uv tool upgrade data-designer" +def test_select_upgrade_command_treats_project_venv_under_uv_tools_as_project() -> None: + command = select_upgrade_command( + environ={}, + python_prefix="/Users/user/projects/uv/tools/my-project/.venv", + ) + + assert command == "uv add --upgrade data-designer" + + def test_select_upgrade_command_detects_pipx_environment() -> None: command = select_upgrade_command( environ={},