Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 94 additions & 0 deletions codeflash/code_utils/version_check.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
"""Version checking utilities for codeflash."""

import sys
import time
from typing import Optional

import requests
from packaging import version

from codeflash.cli_cmds.console import console, logger
from codeflash.version import __version__

# Simple cache to avoid checking too frequently
_version_cache = {"version": None, "timestamp": 0}
_cache_duration = 3600 # 1 hour cache


def get_latest_version_from_pypi() -> Optional[str]:
"""Get the latest version of codeflash from PyPI.

Returns:
The latest version string from PyPI, or None if the request fails.
"""
global _version_cache

# Check cache first
current_time = time.time()
if (_version_cache["version"] is not None and
current_time - _version_cache["timestamp"] < _cache_duration):
return _version_cache["version"]

try:
response = requests.get("https://pypi.org/pypi/codeflash/json", timeout=2)
if response.status_code == 200:
data = response.json()
latest_version = data["info"]["version"]

# Update cache
_version_cache["version"] = latest_version
_version_cache["timestamp"] = current_time

return latest_version
else:
logger.debug(f"Failed to fetch version from PyPI: {response.status_code}")
return None
except requests.RequestException as e:
logger.debug(f"Network error fetching version from PyPI: {e}")
return None
except (KeyError, ValueError) as e:
logger.debug(f"Invalid response format from PyPI: {e}")
return None
except Exception as e:
logger.debug(f"Unexpected error fetching version from PyPI: {e}")
return None


def check_for_newer_minor_version(*, disable_check: bool = False) -> None:
"""Check if a newer minor version is available on PyPI and notify the user.

This function compares the current version with the latest version on PyPI.
If a newer minor version is available, it prints an informational message
suggesting the user upgrade.

Args:
disable_check: If True, skip the version check entirely.
"""
if disable_check:
return

latest_version = get_latest_version_from_pypi()

if not latest_version:
return

try:
current_parsed = version.parse(__version__)
latest_parsed = version.parse(latest_version)

# Check if there's a newer minor version available
# We only notify for minor version updates, not patch updates
if (latest_parsed.major > current_parsed.major or
(latest_parsed.major == current_parsed.major and
latest_parsed.minor > current_parsed.minor)):

console.print(
f"[bold blue]ℹ️ A newer version of Codeflash is available![/bold blue]\n"
f"Current version: {__version__} | Latest version: {latest_version}\n"
f"Consider upgrading for better quality optimizations.",
style="blue"
)

except version.InvalidVersion as e:
logger.debug(f"Invalid version format: {e}")
return
12 changes: 12 additions & 0 deletions codeflash/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from codeflash.cli_cmds.console import paneled_text
from codeflash.code_utils.checkpoint import ask_should_use_checkpoint_get_functions
from codeflash.code_utils.config_parser import parse_config_file
from codeflash.code_utils.version_check import check_for_newer_minor_version
from codeflash.telemetry import posthog_cf
from codeflash.telemetry.sentry import init_sentry

Expand All @@ -21,6 +22,17 @@ def main() -> None:
CODEFLASH_LOGO, panel_args={"title": "https://codeflash.ai", "expand": False}, text_args={"style": "bold gold3"}
)
args = parse_args()

# Check for newer version (skip for version command to avoid confusion)
if not args.version:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With version check 'codeflash --version' we can also show the status if its outdated.

# Check if version check is disabled in config
disable_version_check = False
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure we are going to add this in the configs or existing configs with customers.
We can just have a direct non blocking check.

if args.config_file and Path.exists(args.config_file):
pyproject_config, _ = parse_config_file(args.config_file)
disable_version_check = pyproject_config.get("disable_version_check", False)

check_for_newer_minor_version(disable_check=disable_version_check)

if args.command:
if args.config_file and Path.exists(args.config_file):
pyproject_config, _ = parse_config_file(args.config_file)
Expand Down
1 change: 1 addition & 0 deletions codeflash/version.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
# These version placeholders will be replaced by uv-dynamic-versioning during build.

__version__ = "0.16.6"
227 changes: 227 additions & 0 deletions tests/test_version_check.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
"""Tests for version checking functionality."""

import unittest
from unittest.mock import Mock, patch, MagicMock
from packaging import version

from codeflash.code_utils.version_check import (
get_latest_version_from_pypi,
check_for_newer_minor_version,
_version_cache,
_cache_duration
)


class TestVersionCheck(unittest.TestCase):
"""Test cases for version checking functionality."""

def setUp(self):
"""Reset version cache before each test."""
_version_cache["version"] = None
_version_cache["timestamp"] = 0

def tearDown(self):
"""Clean up after each test."""
_version_cache["version"] = None
_version_cache["timestamp"] = 0

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_success(self, mock_get):
"""Test successful version fetch from PyPI."""
# Mock successful response
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"info": {"version": "1.2.3"}}
mock_get.return_value = mock_response

result = get_latest_version_from_pypi()

self.assertEqual(result, "1.2.3")
mock_get.assert_called_once_with(
"https://pypi.org/pypi/codeflash/json",
timeout=2
)

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_http_error(self, mock_get):
"""Test handling of HTTP error responses."""
# Mock HTTP error response
mock_response = Mock()
mock_response.status_code = 404
mock_get.return_value = mock_response

result = get_latest_version_from_pypi()

self.assertIsNone(result)

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_network_error(self, mock_get):
"""Test handling of network errors."""
# Mock network error
mock_get.side_effect = Exception("Network error")

result = get_latest_version_from_pypi()

self.assertIsNone(result)

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_invalid_response(self, mock_get):
"""Test handling of invalid response format."""
# Mock invalid response format
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"invalid": "format"}
mock_get.return_value = mock_response

result = get_latest_version_from_pypi()

self.assertIsNone(result)

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_caching(self, mock_get):
"""Test that version caching works correctly."""
# Mock successful response
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"info": {"version": "1.2.3"}}
mock_get.return_value = mock_response

# First call should hit the network
result1 = get_latest_version_from_pypi()
self.assertEqual(result1, "1.2.3")
self.assertEqual(mock_get.call_count, 1)

# Second call should use cache
result2 = get_latest_version_from_pypi()
self.assertEqual(result2, "1.2.3")
self.assertEqual(mock_get.call_count, 1) # Still only 1 call

@patch('codeflash.code_utils.version_check.requests.get')
def test_get_latest_version_from_pypi_cache_expiry(self, mock_get):
"""Test that cache expires after the specified duration."""
import time

# Mock successful response
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"info": {"version": "1.2.3"}}
mock_get.return_value = mock_response

# First call
result1 = get_latest_version_from_pypi()
self.assertEqual(result1, "1.2.3")

# Manually expire the cache
_version_cache["timestamp"] = time.time() - _cache_duration - 1

# Second call should hit the network again
result2 = get_latest_version_from_pypi()
self.assertEqual(result2, "1.2.3")
self.assertEqual(mock_get.call_count, 2)

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_newer_available(self, mock_console, mock_get_version):
"""Test warning message when newer minor version is available."""
mock_get_version.return_value = "1.1.0"

check_for_newer_minor_version()

mock_console.print.assert_called_once()
call_args = mock_console.print.call_args[0][0]
self.assertIn("ℹ️ A newer version of Codeflash is available!", call_args)
self.assertIn("Current version: 1.0.0", call_args)
self.assertIn("Latest version: 1.1.0", call_args)

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_newer_major_available(self, mock_console, mock_get_version):
"""Test warning message when newer major version is available."""
mock_get_version.return_value = "2.0.0"

check_for_newer_minor_version()

mock_console.print.assert_called_once()
call_args = mock_console.print.call_args[0][0]
self.assertIn("ℹ️ A newer version of Codeflash is available!", call_args)

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.1.0')
def test_check_for_newer_minor_version_no_newer_available(self, mock_console, mock_get_version):
"""Test no warning when no newer version is available."""
mock_get_version.return_value = "1.0.0"

check_for_newer_minor_version()

mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_patch_update_ignored(self, mock_console, mock_get_version):
"""Test that patch updates don't trigger warnings."""
mock_get_version.return_value = "1.0.1"

check_for_newer_minor_version()

mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_same_version(self, mock_console, mock_get_version):
"""Test no warning when versions are the same."""
mock_get_version.return_value = "1.0.0"

check_for_newer_minor_version()

mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_no_latest_version(self, mock_console, mock_get_version):
"""Test no warning when latest version cannot be fetched."""
mock_get_version.return_value = None

check_for_newer_minor_version()

mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
@patch('codeflash.code_utils.version_check.__version__', '1.0.0')
def test_check_for_newer_minor_version_invalid_version_format(self, mock_console, mock_get_version):
"""Test handling of invalid version format."""
mock_get_version.return_value = "invalid-version"

check_for_newer_minor_version()

mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
def test_check_for_newer_minor_version_disable_check(self, mock_console, mock_get_version):
"""Test that version check is skipped when disable_check is True."""
check_for_newer_minor_version(disable_check=True)

mock_get_version.assert_not_called()
mock_console.print.assert_not_called()

@patch('codeflash.code_utils.version_check.get_latest_version_from_pypi')
@patch('codeflash.code_utils.version_check.console')
def test_check_for_newer_minor_version_disable_check_false(self, mock_console, mock_get_version):
"""Test that version check runs when disable_check is False."""
mock_get_version.return_value = None

check_for_newer_minor_version(disable_check=False)

mock_get_version.assert_called_once()
mock_console.print.assert_not_called()


if __name__ == '__main__':
unittest.main()