diff --git a/bigframes/session/environment.py b/bigframes/session/environment.py index 3ed6ab98cd..940f8deed4 100644 --- a/bigframes/session/environment.py +++ b/bigframes/session/environment.py @@ -1,4 +1,4 @@ -# Copyright 2023 Google LLC +# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,9 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. + import importlib import json import os +import pathlib + +Path = pathlib.Path + # The identifier for GCP VS Code extension # https://cloud.google.com/code/docs/vscode/install @@ -29,40 +34,36 @@ def _is_vscode_extension_installed(extension_id: str) -> bool: """ Checks if a given Visual Studio Code extension is installed. - Args: extension_id: The ID of the extension (e.g., "ms-python.python"). - Returns: True if the extension is installed, False otherwise. """ try: # Determine the user's VS Code extensions directory. - user_home = os.path.expanduser("~") - if os.name == "nt": # Windows - vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions") - elif os.name == "posix": # macOS and Linux - vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions") - else: - raise OSError("Unsupported operating system.") + user_home = Path.home() + vscode_extensions_dir = user_home / ".vscode" / "extensions" # Check if the extensions directory exists. - if os.path.exists(vscode_extensions_dir): - # Iterate through the subdirectories in the extensions directory. - for item in os.listdir(vscode_extensions_dir): - item_path = os.path.join(vscode_extensions_dir, item) - if os.path.isdir(item_path) and item.startswith(extension_id + "-"): - # Check if the folder starts with the extension ID. - # Further check for manifest file, as a more robust check. - manifest_path = os.path.join(item_path, "package.json") - if os.path.exists(manifest_path): - try: - with open(manifest_path, "r", encoding="utf-8") as f: - json.load(f) - return True - except (FileNotFoundError, json.JSONDecodeError): - # Corrupted or incomplete extension, or manifest missing. - pass + if not vscode_extensions_dir.exists(): + return False + + # Iterate through the subdirectories in the extensions directory. + extension_dirs = filter( + lambda p: p.is_dir() and p.name.startswith(extension_id + "-"), + vscode_extensions_dir.iterdir(), + ) + for extension_dir in extension_dirs: + # As a more robust check, the manifest file must exist. + manifest_path = extension_dir / "package.json" + if not manifest_path.exists() or not manifest_path.is_file(): + continue + + # Finally, the manifest file must be a valid json + with open(manifest_path, "r", encoding="utf-8") as f: + json.load(f) + + return True except Exception: pass @@ -72,10 +73,8 @@ def _is_vscode_extension_installed(extension_id: str) -> bool: def _is_package_installed(package_name: str) -> bool: """ Checks if a Python package is installed. - Args: package_name: The name of the package to check (e.g., "requests", "numpy"). - Returns: True if the package is installed, False otherwise. """ diff --git a/tests/unit/session/test_clients.py b/tests/unit/session/test_clients.py index 5d577a52ed..6b0d8583a5 100644 --- a/tests/unit/session/test_clients.py +++ b/tests/unit/session/test_clients.py @@ -13,6 +13,8 @@ # limitations under the License. import os +import pathlib +import tempfile from typing import Optional import unittest.mock as mock @@ -155,6 +157,7 @@ def test_user_agent_not_in_vscode(monkeypatch): monkeypatch_client_constructors(monkeypatch) provider = create_clients_provider() assert_clients_wo_user_agent(provider, "vscode") + assert_clients_wo_user_agent(provider, "googlecloudtools.cloudcode") # We still need to include attribution to bigframes assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") @@ -165,16 +168,48 @@ def test_user_agent_in_vscode(monkeypatch): monkeypatch_client_constructors(monkeypatch) provider = create_clients_provider() assert_clients_w_user_agent(provider, "vscode") + assert_clients_wo_user_agent(provider, "googlecloudtools.cloudcode") # We still need to include attribution to bigframes assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") +@mock.patch.dict(os.environ, {"VSCODE_PID": "12345"}, clear=True) +def test_user_agent_in_vscode_w_extension(monkeypatch): + monkeypatch_client_constructors(monkeypatch) + + with tempfile.TemporaryDirectory() as tmpdir: + user_home = pathlib.Path(tmpdir) + extension_dir = ( + user_home / ".vscode" / "extensions" / "googlecloudtools.cloudcode-0.12" + ) + extension_config = extension_dir / "package.json" + + # originally extension config does not exist + assert not extension_config.exists() + + # simulate extension installation by creating extension config on disk + extension_dir.mkdir(parents=True) + with open(extension_config, "w") as f: + f.write("{}") + + with mock.patch("pathlib.Path.home", return_value=user_home): + provider = create_clients_provider() + assert_clients_w_user_agent(provider, "vscode") + assert_clients_w_user_agent(provider, "googlecloudtools.cloudcode") + + # We still need to include attribution to bigframes + assert_clients_w_user_agent( + provider, f"bigframes/{bigframes.version.__version__}" + ) + + @mock.patch.dict(os.environ, {}, clear=True) def test_user_agent_not_in_jupyter(monkeypatch): monkeypatch_client_constructors(monkeypatch) provider = create_clients_provider() assert_clients_wo_user_agent(provider, "jupyter") + assert_clients_wo_user_agent(provider, "bigquery_jupyter_plugin") # We still need to include attribution to bigframes assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") @@ -185,6 +220,37 @@ def test_user_agent_in_jupyter(monkeypatch): monkeypatch_client_constructors(monkeypatch) provider = create_clients_provider() assert_clients_w_user_agent(provider, "jupyter") + assert_clients_wo_user_agent(provider, "bigquery_jupyter_plugin") # We still need to include attribution to bigframes assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}") + + +@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "12345"}, clear=True) +def test_user_agent_in_jupyter_with_plugin(monkeypatch): + monkeypatch_client_constructors(monkeypatch) + + def custom_import_module_side_effect(name, package=None): + if name == "bigquery_jupyter_plugin": + return mock.MagicMock() + else: + import importlib + + return importlib.import_module(name, package) + + assert isinstance( + custom_import_module_side_effect("bigquery_jupyter_plugin"), mock.MagicMock + ) + assert custom_import_module_side_effect("bigframes") is bigframes + + with mock.patch( + "importlib.import_module", side_effect=custom_import_module_side_effect + ): + provider = create_clients_provider() + assert_clients_w_user_agent(provider, "jupyter") + assert_clients_w_user_agent(provider, "bigquery_jupyter_plugin") + + # We still need to include attribution to bigframes + assert_clients_w_user_agent( + provider, f"bigframes/{bigframes.version.__version__}" + )