diff --git a/google/cloud/dataproc_spark_connect/environment.py b/google/cloud/dataproc_spark_connect/environment.py index c74e81e..e19dd97 100644 --- a/google/cloud/dataproc_spark_connect/environment.py +++ b/google/cloud/dataproc_spark_connect/environment.py @@ -17,6 +17,11 @@ from typing import Callable, Tuple, List +def is_antigravity() -> bool: + """True if running inside the Antigravity environment.""" + return "antigravity" in os.getenv("__CFBundleIdentifier", "").lower() + + def is_vscode() -> bool: """True if running inside VS Code at all.""" return os.getenv("VSCODE_PID") is not None @@ -148,12 +153,13 @@ def get_client_environment_label() -> str: 10. Google Cloud Shell ("cloud-shell") 11. Hex ("hex") 12. Polynote ("polynote") - 13. VS Code ("vscode") - 14. JetBrains IDE ("jetbrains") - 15. Spyder ("spyder") - 16. Eclipse ("eclipse") - 17. Jupyter ("jupyter") - 18. Unknown ("unknown") + 13. Antigravity ("antigravity") + 14. VS Code ("vscode") + 15. JetBrains IDE ("jetbrains") + 16. Spyder ("spyder") + 17. Eclipse ("eclipse") + 18. Jupyter ("jupyter") + 19. Unknown ("unknown") """ checks: List[Tuple[Callable[[], bool], str]] = [ (is_colab_enterprise, "colab-enterprise"), @@ -168,6 +174,7 @@ def get_client_environment_label() -> str: (is_cloud_shell, "cloud-shell"), (is_hex, "hex"), (is_polynote, "polynote"), + (is_antigravity, "antigravity"), (is_vscode, "vscode"), (is_jetbrains_ide, "jetbrains"), (is_spyder, "spyder"), diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index d3bd5f0..a64387a 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -30,6 +30,22 @@ def tearDown(self): os.environ.clear() os.environ.update(self.original_environ) + def test_is_antigravity_true(self): + os.environ["__CFBundleIdentifier"] = "com.google.antigravity" + self.assertTrue(environment.is_antigravity()) + + def test_is_antigravity_true_case_insensitive(self): + os.environ["__CFBundleIdentifier"] = "AntiGravity-App" + self.assertTrue(environment.is_antigravity()) + + def test_is_antigravity_false_not_present(self): + os.environ.pop("__CFBundleIdentifier", None) + self.assertFalse(environment.is_antigravity()) + + def test_is_antigravity_false_different_value(self): + os.environ["__CFBundleIdentifier"] = "com.example.app" + self.assertFalse(environment.is_antigravity()) + def test_is_vscode_true(self): os.environ["VSCODE_PID"] = "12345" self.assertTrue(environment.is_vscode()) @@ -231,6 +247,10 @@ def test_is_jetbrains_ide_false_env_var_not_jetbrains(self): "google.cloud.dataproc_spark_connect.environment.is_polynote", return_value=False, ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_antigravity", + return_value=False, + ) @mock.patch( "google.cloud.dataproc_spark_connect.environment.is_vscode", return_value=False, @@ -371,6 +391,68 @@ def test_get_client_environment_label_kaggle(self, *mocks): "google.cloud.dataproc_spark_connect.environment.is_polynote", return_value=False, ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_antigravity", + return_value=True, + ) + def test_get_client_environment_label_antigravity(self, *mocks): + self.assertEqual( + environment.get_client_environment_label(), + "antigravity", + ) + + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_colab_enterprise", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_colab", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_workbench", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_kaggle", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_sagemaker", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_databricks", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_deepnote", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_datalore", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_codespaces", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_cloud_shell", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_hex", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_polynote", + return_value=False, + ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_antigravity", + return_value=False, + ) @mock.patch( "google.cloud.dataproc_spark_connect.environment.is_vscode", return_value=True, @@ -429,6 +511,10 @@ def test_get_client_environment_label_vscode(self, *mocks): "google.cloud.dataproc_spark_connect.environment.is_polynote", return_value=False, ) + @mock.patch( + "google.cloud.dataproc_spark_connect.environment.is_antigravity", + return_value=False, + ) @mock.patch( "google.cloud.dataproc_spark_connect.environment.is_vscode", return_value=False,