feat: Detect Antigravity IDE for usage tracking (#183)

Author: virrrat

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"),

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,