Feat: Introduce ExtendedGcsFileSystem for Zonal Bucket gRPC Read Path

.github/workflows/ci.yml

@@ -35,13 +35,22 @@ jobs:
       - name: install
         run: |
           pip install -e .
-      - name: Run tests
+      - name: Run Standard Tests
         run: |
           export GOOGLE_APPLICATION_CREDENTIALS=$(pwd)/gcsfs/tests/fake-secret.json
           pytest -vv -s \
           --log-format="%(asctime)s %(levelname)s %(message)s" \
           --log-date-format="%H:%M:%S" \
-          gcsfs/
+          gcsfs/ \
+          --ignore=gcsfs/tests/test_extended_gcsfs.py
+      - name: Run Extended Tests
+        run: |
+          export GOOGLE_APPLICATION_CREDENTIALS=$(pwd)/gcsfs/tests/fake-secret.json
+          export GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT="true"
+          pytest -vv -s \
+            --log-format="%(asctime)s %(levelname)s %(message)s" \
+            --log-date-format="%H:%M:%S" \
+            gcsfs/tests/test_extended_gcsfs.py
 
   lint:
     name: lint

gcsfs/__init__.py

@@ -1,10 +1,27 @@
+import logging
+import os
+
 from ._version import get_versions
 
+logger = logging.getLogger(__name__)
 __version__ = get_versions()["version"]
 del get_versions
 from .core import GCSFileSystem
 from .mapping import GCSMap
 
+if os.getenv("GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT", "false").lower() in ("true", "1"):
+    try:
+        from .extended_gcsfs import ExtendedGcsFileSystem as GCSFileSystem
+
+        logger.info(
+            "gcsfs experimental features enabled via GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT."
+        )
+    except ImportError as e:
+        logger.warning(
+            f"GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT is set, but failed to import experimental features: {e}"
+        )
+        # Fallback to core GCSFileSystem, do not register here
+
 __all__ = ["GCSFileSystem", "GCSMap"]
 
 from . import _version

gcsfs/core.py

@@ -282,22 +282,6 @@ class GCSFileSystem(asyn.AsyncFileSystem):
     protocol = "gs", "gcs"
     async_impl = True
 
-    def __new__(cls, *args, **kwargs):
-        """
-        Factory to return a ExtendedGcsFileSystem instance if the experimental
-        flag is enabled.
-        """
-        experimental_multi_bucket_support = os.environ.get(
-            "GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT", "false"
-        ).lower() in ("true", "1")
-
-        if experimental_multi_bucket_support:
-            from .extended_gcsfs import ExtendedGcsFileSystem
-
-            return object.__new__(ExtendedGcsFileSystem)
-        else:
-            return object.__new__(cls)
-
     def __init__(
         self,
         project=DEFAULT_PROJECT,

gcsfs/tests/test_core.py

@@ -4,7 +4,6 @@
 from datetime import datetime, timezone
 from itertools import chain
 from unittest import mock
-from unittest.mock import patch
 from urllib.parse import parse_qs, unquote, urlparse
 from uuid import uuid4
 
@@ -19,7 +18,6 @@
 from gcsfs import __version__ as version
 from gcsfs.core import GCSFileSystem, quote
 from gcsfs.credentials import GoogleCredentials
-from gcsfs.extended_gcsfs import ExtendedGcsFileSystem
 from gcsfs.tests.conftest import a, allfiles, b, csv_files, files, text_files
 from gcsfs.tests.utils import tempdir, tmpfile
 
@@ -1733,27 +1731,3 @@ def test_near_find(gcs):
 def test_get_error(gcs):
     with pytest.raises(FileNotFoundError):
         gcs.get_file(f"{TEST_BUCKET}/doesnotexist", "other")
-
-
-def test_gcs_filesystem_when_experimental_zonal_toggle_is_not_present(gcs_factory):
-    gcs = gcs_factory()
-
-    assert isinstance(
-        gcs, gcsfs.GCSFileSystem
-    ), "Expected File system instance to be GCSFileSystem"
-    assert not isinstance(
-        gcs, ExtendedGcsFileSystem
-    ), "Expected File system instance to be GCSFileSystem"
-
-
-def test_extended_gcs_filesystem_when_experimental_zonal_toggle_is_true(gcs_factory):
-    try:
-        with patch("google.auth.default", return_value=(None, "fake-project")):
-            os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"] = "true"
-            gcs = gcs_factory()
-
-            assert isinstance(
-                gcs, ExtendedGcsFileSystem
-            ), "Expected File system instance to be ExtendedGcsFileSystem"
-    finally:
-        del os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"]

gcsfs/tests/test_init.py

@@ -0,0 +1,67 @@
+import os
+import sys
+
+
+class TestConditionalImport:
+    def setup_method(self, method):
+        """Setup for each test method."""
+        self.original_env = os.environ.get("GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT")
+
+        # Snapshot original gcsfs modules
+        self.original_modules = {
+            name: mod for name, mod in sys.modules.items() if name.startswith("gcsfs")
+        }
+
+        # Unload gcsfs modules to force re-import during the test
+        modules_to_remove = list(self.original_modules.keys())
+        for name in modules_to_remove:
+            if name in sys.modules:
+                del sys.modules[name]
+
+    def teardown_method(self, method):
+        """Teardown after each test method."""
+        # Reset environment variable to its original state
+        if self.original_env is not None:
+            os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"] = self.original_env
+        elif "GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT" in os.environ:
+            del os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"]
+
+        # Clear any gcsfs modules loaded/modified during this test
+        modules_to_remove = [name for name in sys.modules if name.startswith("gcsfs")]
+        for name in modules_to_remove:
+            if name in sys.modules:
+                del sys.modules[name]
+
+        # Restore the original gcsfs modules from the snapshot to avoid side effect
+        # affecting other tests
+        sys.modules.update(self.original_modules)
+
+    def test_experimental_env_unset(self):
+        """
+        Tests gcsfs.GCSFileSystem is core.GCSFileSystem when
+        GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT is NOT set.
+        """
+        if "GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT" in os.environ:
+            del os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"]
+
+        import gcsfs
+
+        assert (
+            gcsfs.GCSFileSystem is gcsfs.core.GCSFileSystem
+        ), "Should be core.GCSFileSystem"
+        assert not hasattr(
+            gcsfs, "ExtendedGcsFileSystem"
+        ), "ExtendedGcsFileSystem should not be imported directly on gcsfs"
+
+    def test_experimental_env_set(self):
+        """
+        Tests gcsfs.GCSFileSystem is extended_gcsfs.ExtendedGcsFileSystem when
+        GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT IS set.
+        """
+        os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"] = "true"
+
+        import gcsfs
+
+        assert (
+            gcsfs.GCSFileSystem is gcsfs.extended_gcsfs.ExtendedGcsFileSystem
+        ), "Should be ExtendedGcsFileSystem"