Merge pull request #26901 from NeilGirdhar:etils

PiperOrigin-RevId: 733466732

jax/_src/compilation_cache_interface.py

@@ -15,8 +15,8 @@
 from __future__ import annotations
 
 import abc
+import pathlib
 
-from jax._src import path as pathlib
 from jax._src import util
 
 

jax/_src/path.py

@@ -12,35 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Protocol
 import logging
+import os
 import pathlib
-import importlib.util
 
 __all__ = ["Path"]
-
 logger = logging.getLogger(__name__)
+epath_installed: bool
+
+class PathProtocol(Protocol):
+  """A factory that creates a PurePath."""
+  def __call__(self, *pathsegments: str | os.PathLike) -> pathlib.Path:
+    ...
+
+Path: PathProtocol
 
 # If etils.epath (aka etils[epath] to pip) is present, we prefer it because it
 # can read and write to, e.g., GCS buckets. Otherwise we use the builtin
 # pathlib and can only read/write to the local filesystem.
-epath_installed = bool(
-    importlib.util.find_spec("etils") and
-    importlib.util.find_spec("etils.epath")
-)
-if epath_installed:
-  logger.debug("etils.epath found. Using etils.epath for file I/O.")
-
-  def __dir__():
-    return ["Path"]
-
-  def __getattr__(name):
-    if name != "Path":
-      raise AttributeError(f"module '{__name__}' has no attribute '{name}")
-
-    global Path
-    from etils import epath
-    Path = epath.Path
-    return Path
-else:
+try:
+  from etils import epath  # type: ignore
+except ImportError:
   logger.debug("etils.epath was not found. Using pathlib for file I/O.")
   Path = pathlib.Path
+  epath_installed = False
+else:
+  logger.debug("etils.epath found. Using etils.epath for file I/O.")
+  # Ultimately, epath.Path implements pathlib.Path.  See:
+  # https://github.com/google/etils/blob/2083f3d932a88d8a135ef57112cd1f9aff5d559e/etils/epath/abstract_path.py#L47
+  Path = epath.Path
+  epath_installed = True