Skip to content
84 changes: 84 additions & 0 deletions codeflash/code_utils/version_check.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
"""Version checking utilities for codeflash."""

from __future__ import annotations

import time

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() -> str | None:
"""Get the latest version of codeflash from PyPI.

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

"""
# 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
logger.debug(f"Failed to fetch version from PyPI: {response.status_code}")
return None # noqa: TRY300
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() -> 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.
"""
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
):
Comment on lines +72 to +74
Copy link
Contributor

Choose a reason for hiding this comment

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

personally, I prefer checking if latest_version < __version__ (major or minor)
@Saga4 what do you think ?

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",
)
Comment on lines +75 to +80
Copy link
Contributor

Choose a reason for hiding this comment

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

that's a lot of info, maybe just a simple one line message would do.

also I'm not a big fan of for better quality optimizations,
the message should say something like "Things could behave unexpectedly"

also use the warning log level for that


except version.InvalidVersion as e:
logger.debug(f"Invalid version format: {e}")
return
8 changes: 6 additions & 2 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,12 +22,15 @@ 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 for all commands
check_for_newer_minor_version()

if args.command:
disable_telemetry = False
if args.config_file and Path.exists(args.config_file):
pyproject_config, _ = parse_config_file(args.config_file)
disable_telemetry = pyproject_config.get("disable_telemetry", False)
else:
disable_telemetry = False
init_sentry(not disable_telemetry, exclude_errors=True)
posthog_cf.initialize_posthog(not disable_telemetry)
args.func()
Expand Down
208 changes: 208 additions & 0 deletions tests/test_version_check.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,208 @@
"""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()



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