diff --git a/codeflash/code_utils/version_check.py b/codeflash/code_utils/version_check.py new file mode 100644 index 00000000..0eba9e09 --- /dev/null +++ b/codeflash/code_utils/version_check.py @@ -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 + ): + 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 diff --git a/codeflash/main.py b/codeflash/main.py index 650bdbd6..27376aef 100644 --- a/codeflash/main.py +++ b/codeflash/main.py @@ -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 @@ -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() diff --git a/tests/test_version_check.py b/tests/test_version_check.py new file mode 100644 index 00000000..25c4c7bc --- /dev/null +++ b/tests/test_version_check.py @@ -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() \ No newline at end of file