From b8b330a25df2ed394152a092c73f8830857d0987 Mon Sep 17 00:00:00 2001 From: Daniel Ng Date: Thu, 29 Oct 2026 10:53:05 -0700 Subject: [PATCH] Auto-select gcs_grpc driver for GCS Zonal Rapid buckets PiperOrigin-RevId: 1009806965 --- CHANGELOG.md | 8 ++- etils/__init__.py | 2 +- etils/ecolab/auto_display_utils.py | 12 ++-- etils/ecolab/cell_autoreload.py | 2 +- etils/ecolab/inplace_reload.py | 2 +- etils/ecolab/lazy_imports.py | 96 +++++++++++++++--------------- etils/ecolab/lazy_imports_test.py | 10 ++-- etils/ecolab/patch_utils.py | 14 ++--- etils/edc/cast_utils.py | 2 +- etils/edc/cast_utils_test.py | 4 +- etils/edc/context.py | 2 +- etils/edc/dataclass_utils_test.py | 2 +- etils/edc/field_utils_test.py | 8 +-- etils/edc/frozen_utils.py | 4 +- etils/edc/frozen_utils_test.py | 34 +++++------ etils/edc/helpers.py | 2 +- etils/epath/abstract_path.py | 15 ++--- etils/epath/backend.py | 14 +++-- etils/epath/backend_test.py | 9 +++ etils/epath/flags.py | 2 +- etils/epath/gpath.py | 30 +++++++++- etils/epath/gpath_test.py | 58 ++++++++++++++++-- etils/epath/register.py | 8 +-- etils/epath/resource_utils.py | 16 ++--- etils/epath/testing.py | 2 +- pyproject.toml | 2 +- 26 files changed, 224 insertions(+), 136 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b5c997d..963928a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ Changelog follow https://keepachangelog.com/ format. ## [Unreleased] +## [1.15.0] - 2026-10-29 + +* `epath`: + * Support GCS Zonal Rapid (`gs://`) buckets via `fsspec` + (`gcsfs>=2026.4.0`). * `edc`: * [Fix] `edc.ContextVar`: `default_factory` is now lazily resolved per-context, so each thread / coroutine gets its own independent @@ -537,7 +542,8 @@ Changelog follow https://keepachangelog.com/ format. -[Unreleased]: https://github.com/google/etils/compare/v1.14.0...HEAD +[Unreleased]: https://github.com/google/etils/compare/v1.15.0...HEAD +[1.15.0]: https://github.com/google/etils/compare/v1.14.0...v1.15.0 [1.14.0]: https://github.com/google/etils/compare/v1.13.0...v1.14.0 [1.13.0]: https://github.com/google/etils/compare/v1.12.2...v1.13.0 [1.12.2]: https://github.com/google/etils/compare/v1.12.1...v1.12.2 diff --git a/etils/__init__.py b/etils/__init__.py index b847b0d8..b89f9cd8 100644 --- a/etils/__init__.py +++ b/etils/__init__.py @@ -16,7 +16,7 @@ # A new PyPI release will be pushed everytime `__version__` is increased # When changing this, also update the CHANGELOG.md -__version__ = '1.14.0' +__version__ = '1.15.0' # Do NOT add anything to this file. This is left empty on purpose. # Users should import subprojects directly. diff --git a/etils/ecolab/auto_display_utils.py b/etils/ecolab/auto_display_utils.py index b2eff61a..1e1db682 100644 --- a/etils/ecolab/auto_display_utils.py +++ b/etils/ecolab/auto_display_utils.py @@ -225,7 +225,7 @@ def _reraise_error(fn: _T) -> _T: @functools.wraps(fn) # pyrefly: ignore[bad-argument-type] def decorated(self, node: ast.AST): try: - return fn(self, node) # pytype: disable=wrong-arg-types + return fn(self, node) # pyrefly: ignore[not-callable] except Exception as e: # pylint: disable=broad-exception-caught code = '\n'.join(self.lines_recorder.last_lines) print(f'Error for code:\n-----\n{code}\n-----') @@ -260,7 +260,7 @@ def _maybe_display( if line_info.print_line: fn_kwargs.append(ast.keyword('line_code', _unparse_line(node))) - node.value = ast.Call( # pytype: disable=wrong-arg-types + node.value = ast.Call( func=_parse_expr('ecolab.auto_display_utils._display_and_return'), # pyrefly: ignore[bad-argument-type] args=[node.value], keywords=fn_kwargs, @@ -289,7 +289,7 @@ def _is_alias_stmt(self, node: ast.AST) -> bool: def visit_Assert(self, node: ast.Assert) -> None: # pylint: disable=invalid-name # Wrap assert so the `node.value` match the expected API node = _WrapAssertNode(node) - node = self._maybe_display(node) # pytype: disable=wrong-arg-types + node = self._maybe_display(node) # pyrefly: ignore[bad-argument-type, bad-assignment] assert isinstance(node, _WrapAssertNode) node = node._node # Unwrap # pylint: disable=protected-access return node # pyrefly: ignore[bad-return] @@ -351,12 +351,12 @@ def _has_trailing_semicolon( ) # Extract the lines of the statement - line_num = node.end_lineno - 1 # pytype: disable=attribute-error + line_num = node.end_lineno - 1 # pyrefly: ignore[missing-attribute] last_line = code_lines[line_num] # lineno starts at `1` # `node.end_col_offset` is in bytes, so UTF-8 characters count 3. last_part_of_line = last_line.encode('utf-8') - last_part_of_line = last_part_of_line[node.end_col_offset :] # pytype: disable=attribute-error + last_part_of_line = last_part_of_line[node.end_col_offset :] last_part_of_line = last_part_of_line.decode('utf-8') # Check if the last character is a `;` token @@ -386,7 +386,7 @@ def _detect_trailing_regex() -> re.Pattern[str]: # * `; a; b` # * `; a=1` - available_letters = ''.join(sorted(_Options.all_letters)) # pytype: disable=wrong-arg-types + available_letters = ''.join(sorted(_Options.all_letters)) return re.compile( ' *; *' # Trailing `;` (surrounded by spaces) f'(?P[{available_letters}]*)?' # Optionally a `option` letter diff --git a/etils/ecolab/cell_autoreload.py b/etils/ecolab/cell_autoreload.py index 334cd0d0..3f98298e 100644 --- a/etils/ecolab/cell_autoreload.py +++ b/etils/ecolab/cell_autoreload.py @@ -158,7 +158,7 @@ def _pre_run_cell_maybe_reload( 'reload_recursive': False, 'collapse_prefix': f'Autoreload ({len(modules_to_reload)} modules): ', } - with adhoc_imports.adhoc(**adhoc_kwargs): # pyrefly: ignore[bad-argument-type] + with adhoc_imports.adhoc(**adhoc_kwargs): for module in modules_to_reload: importlib.import_module(module) diff --git a/etils/ecolab/inplace_reload.py b/etils/ecolab/inplace_reload.py index 931c6608..2948c3ad 100644 --- a/etils/ecolab/inplace_reload.py +++ b/etils/ecolab/inplace_reload.py @@ -151,7 +151,7 @@ def update(self, old, new): case types.FunctionType(), types.FunctionType(): return self._update_function(old, new) case types.MethodType(), types.MethodType(): - return self._update_function(old.__func__, new.__func__) # pytype: disable=wrong-arg-types + return self._update_function(old.__func__, new.__func__) case property(), property(): return self._update_property(old, new) diff --git a/etils/ecolab/lazy_imports.py b/etils/ecolab/lazy_imports.py index 90e934cc..111d5b17 100644 --- a/etils/ecolab/lazy_imports.py +++ b/etils/ecolab/lazy_imports.py @@ -117,7 +117,7 @@ def print_current_imports() -> None: import threading import time import timeit - import tomllib # pytype: disable=import-error + import tomllib import traceback import typing # Note we do not import `Any`, `TypeVar`,... import types @@ -132,72 +132,71 @@ def print_current_imports() -> None: with _builder.replace_imports(is_std=False): # pylint: disable=g-import-not-at-top,unused-import,reimported - # pytype: disable=import-error # ====== Etils ====== from etils import array_types from etils import ecolab - from etils import edc + from etils import edc # pyrefly: ignore[missing-module-attribute] from etils import enp from etils import epath from etils import epy from etils import etqdm from etils import etree - from etils import exm + from etils import exm # pyrefly: ignore[missing-module-attribute] from etils import g3_utils from etils.ecolab import lazy_imports # ====== Common third party ====== from absl import app from absl import flags - import apache_beam as beam - import bagz - import chex - import dataclass_array as dca + import apache_beam as beam # pyrefly: ignore[missing-import] + import bagz # pyrefly: ignore[missing-import] + import chex # pyrefly: ignore[missing-import] + import dataclass_array as dca # pyrefly: ignore[missing-import] import einops - import fiddle as fdl - import flask - import flax - from flax import linen as nn - from flax import nnx - import functorch - import gin - import grain.python as grain - import graphviz - import imageio + import fiddle as fdl # pyrefly: ignore[missing-import] + import flask # pyrefly: ignore[missing-import] + import flax # pyrefly: ignore[missing-import] + from flax import linen as nn # pyrefly: ignore[missing-import] + from flax import nnx # pyrefly: ignore[missing-import] + import functorch # pyrefly: ignore[missing-import] + import gin # pyrefly: ignore[missing-import] + import grain.python as grain # pyrefly: ignore[missing-import] + import graphviz # pyrefly: ignore[missing-import] + import imageio # pyrefly: ignore[missing-import] import immutabledict # Even though `import ipywidgets as widgets` is the common alias, widgets # is likely too ambiguous. import ipywidgets - import jax - from jax import numpy as jnp - import jaxtyping - import lark + import jax # pyrefly: ignore[missing-import] + from jax import numpy as jnp # pyrefly: ignore[missing-import] + import jaxtyping # pyrefly: ignore[missing-import] + import lark # pyrefly: ignore[missing-import] import matplotlib import matplotlib as mpl # Standard alias from matplotlib import pyplot as plt - import mcp + import mcp # pyrefly: ignore[missing-import] import mediapy as media - import ml_collections - import networkx as nx + import ml_collections # pyrefly: ignore[missing-import] + import networkx as nx # pyrefly: ignore[missing-source-for-stubs] import numpy as np - import optax - import orbax - from orbax import checkpoint as ocp - from orbax.checkpoint.experimental import v1 as ocp_v1 - from orbax import export as oex - import pandas as pd + import optax # pyrefly: ignore[missing-import] + import orbax # pyrefly: ignore[missing-import] + from orbax import checkpoint as ocp # pyrefly: ignore[missing-import] + from orbax.checkpoint.experimental import v1 as ocp_v1 # pyrefly: ignore[missing-import] + from orbax import export as oex # pyrefly: ignore[missing-import] + import pandas as pd # pyrefly: ignore[missing-import] import PIL from PIL import Image # Common alias - import pycolmap - import scipy - import seaborn as sns - import sklearn - import tensorflow as tf - import tensorflow.experimental.numpy as tnp - import tensorflow_datasets as tfds - import torch + import pycolmap # pyrefly: ignore[missing-import] + import scipy # pyrefly: ignore[missing-import] + import seaborn as sns # pyrefly: ignore[missing-source-for-stubs] + import sklearn # pyrefly: ignore[missing-import] + import tensorflow as tf # pyrefly: ignore[missing-source-for-stubs] + import tensorflow.experimental.numpy as tnp # pyrefly: ignore[missing-import] + import tensorflow_datasets as tfds # pyrefly: ignore[missing-import] + import torch # pyrefly: ignore[missing-import] # from torch import nn # Collision with flax.linen - import torchtext - import torchvision + import torchtext # pyrefly: ignore[missing-import] + import torchvision # pyrefly: ignore[missing-import] import tqdm # tqdm import also trigger additional imports. # TODO(epot): Currently pylance might not infer `tqdm.auto` match @@ -206,18 +205,17 @@ def print_current_imports() -> None: tqdm.auto # pylint: disable=pointless-statement tqdm.notebook # pylint: disable=pointless-statement import tree - import typeguard + import typeguard # pyrefly: ignore[missing-import] import typing_extensions - import plotly - from plotly import express as px - from plotly import graph_objects as go + import plotly # pyrefly: ignore[missing-import] + from plotly import express as px # pyrefly: ignore[missing-import] + from plotly import graph_objects as go # pyrefly: ignore[missing-import] from pydantic import v1 as pydantic import requests - import sunds - import visu3d as v3d - from xmanager.contrib import flow as xmflow + import sunds # pyrefly: ignore[missing-import] + import visu3d as v3d # pyrefly: ignore[missing-import] + from xmanager.contrib import flow as xmflow # pyrefly: ignore[missing-import] from xmanager import xm - # pytype: enable=import-error # pylint: enable=g-import-not-at-top,unused-import,reimported diff --git a/etils/ecolab/lazy_imports_test.py b/etils/ecolab/lazy_imports_test.py index 7a49f693..076c90a7 100644 --- a/etils/ecolab/lazy_imports_test.py +++ b/etils/ecolab/lazy_imports_test.py @@ -56,9 +56,9 @@ def test_lazy_imports(): # lazy_imports can be imported but do not appear in `print_current_imports` from etils.ecolab.lazy_imports import lazy_imports # pylint: disable=g-import-not-at-top - assert not lazy_imports._etils_state.module_loaded # pytype: disable=module-attr + assert not lazy_imports._etils_state.module_loaded # pyrefly: ignore[missing-attribute] _ = lazy_imports.print_current_imports # Trigger import - assert lazy_imports._etils_state.module_loaded # pytype: disable=module-attr + assert lazy_imports._etils_state.module_loaded # pyrefly: ignore[missing-attribute] del lazy_imports assert ecolab.lazy_utils.current_import_statements( @@ -75,16 +75,16 @@ def test_lazy_imports(): def test_lazy_imports_mutate(): from etils.lazy_imports import enp # pylint: disable=g-import-not-at-top - assert not enp._etils_state.module_loaded # pytype: disable=module-attr + assert not enp._etils_state.module_loaded # pyrefly: ignore[missing-attribute] _ = enp.NpModule - assert enp._etils_state.module_loaded # pytype: disable=module-attr + assert enp._etils_state.module_loaded # pyrefly: ignore[missing-attribute] enp.some_value = 123 # pyrefly: ignore[missing-attribute] # Mutating lazy import is propagated to the shared instance. from etils import enp as enp2 # pylint: disable=g-import-not-at-top - assert enp2.some_value == 123 # pytype: disable=module-attr + assert enp2.some_value == 123 # pyrefly: ignore[missing-attribute] def test_lazy_imports_built_in(): diff --git a/etils/ecolab/patch_utils.py b/etils/ecolab/patch_utils.py index d438df2a..8ab88463 100644 --- a/etils/ecolab/patch_utils.py +++ b/etils/ecolab/patch_utils.py @@ -22,10 +22,8 @@ def set_verbose() -> None: """Log stderr & `absl.logging` in Colab (filtered by default).""" # pylint: disable=g-import-not-at-top - # pytype: disable=import-error from absl import logging - from colabtools import googlelog - # pytype: enable=import-error + from colabtools import googlelog # pyrefly: ignore[missing-import] # pylint: enable=g-import-not-at-top logging.set_verbosity(logging.INFO) @@ -51,13 +49,11 @@ def patch_graphviz() -> None: Calling this function fix the behavior. """ # pylint: disable=g-import-not-at-top - # pytype: disable=import-error - from colabtools import proto - from colabtools import publish - from colabtools import stubby + from colabtools import proto # pyrefly: ignore[missing-import] + from colabtools import publish # pyrefly: ignore[missing-import] + from colabtools import stubby # pyrefly: ignore[missing-import] - import graphviz - # pytype: enable=import-error + import graphviz # pyrefly: ignore[missing-import] # pylint: enable=g-import-not-at-top request_proto_cls = proto.GetProtoClass('graphviz_server.RenderRequest') diff --git a/etils/edc/cast_utils.py b/etils/edc/cast_utils.py index 465bd618..7b1ffd3c 100644 --- a/etils/edc/cast_utils.py +++ b/etils/edc/cast_utils.py @@ -41,7 +41,7 @@ class AutoCast(metaclass=_AutoCastMeta): pass else: - AutoCast = Annotated[_T, _IS_NORMALIZED] # pytype: disable=invalid-typevar + AutoCast = Annotated[_T, _IS_NORMALIZED] def make_auto_cast_descriptor( diff --git a/etils/edc/cast_utils_test.py b/etils/edc/cast_utils_test.py index 1cc3a597..86e9db26 100644 --- a/etils/edc/cast_utils_test.py +++ b/etils/edc/cast_utils_test.py @@ -41,9 +41,9 @@ class C(B): class D(C): d: edc.AutoCast[ int - ] = '888' # Default arg # pytype: disable=annotation-type-mismatch + ] = '888' # Default arg - d = D(c='123', c_non_autocast='456') # pytype: disable=wrong-arg-types + d = D(c='123', c_non_autocast='456') # pyrefly: ignore[bad-argument-type] assert d.d == 888 assert d.c == 123 diff --git a/etils/edc/context.py b/etils/edc/context.py index 57d996c8..9dd7abc0 100644 --- a/etils/edc/context.py +++ b/etils/edc/context.py @@ -44,7 +44,7 @@ class ContextVar(metaclass=_ContextVarMeta): pass else: - ContextVar = Annotated[_T, _IS_CONTEXTVAR] # pytype: disable=invalid-typevar + ContextVar = Annotated[_T, _IS_CONTEXTVAR] def make_contextvar_descriptor( diff --git a/etils/edc/dataclass_utils_test.py b/etils/edc/dataclass_utils_test.py index 14eeae85..a302d2f5 100644 --- a/etils/edc/dataclass_utils_test.py +++ b/etils/edc/dataclass_utils_test.py @@ -57,7 +57,7 @@ def replace(self) -> int: # Custom replace function def test_replace(): obj = object() x = A(y=obj) - y = x.replace(x=123) # pytype: disable=attribute-error + y = x.replace(x=123) # pyrefly: ignore[missing-attribute] assert x == A(y=obj) assert y == A(x=123, y=obj) assert x.y is y.y diff --git a/etils/edc/field_utils_test.py b/etils/edc/field_utils_test.py index 613b7175..14c6bf02 100644 --- a/etils/edc/field_utils_test.py +++ b/etils/edc/field_utils_test.py @@ -27,7 +27,7 @@ def test_field_no_op(frozen: bool): @dataclasses.dataclass(frozen=frozen) class A: - x: Any = edc.field() # pytype: disable=annotation-type-mismatch + x: Any = edc.field() # No argument with pytest.raises(TypeError, match=r'__init__\(\) missing 1 required'): @@ -51,7 +51,7 @@ def test_field_multi_instance(frozen: bool): @dataclasses.dataclass(frozen=frozen) class A: - x: Any = edc.field(validate=str) # pytype: disable=annotation-type-mismatch + x: Any = edc.field(validate=str) a0 = A(123) assert a0.x == '123' @@ -95,10 +95,10 @@ def test_field_validate( eq=eq, order=order, unsafe_hash=unsafe_hash, - frozen=frozen, # pytype: disable=not-supported-yet + frozen=frozen, ) class A: - x: Any = edc.field(validate=str, **default_kwargs) # pytype: disable=annotation-type-mismatch + x: Any = edc.field(validate=str, **default_kwargs) # No argument if default_kwargs: diff --git a/etils/edc/frozen_utils.py b/etils/edc/frozen_utils.py index a9c00a6c..87a7a241 100644 --- a/etils/edc/frozen_utils.py +++ b/etils/edc/frozen_utils.py @@ -99,7 +99,7 @@ def __repr__(self) -> str: def tree_flatten(self) -> tuple[list[Any], Any]: """`jax.tree_utils` support.""" - import jax # pylint: disable=g-import-not-at-top # pytype: disable=import-error + import jax # pylint: disable=g-import-not-at-top # pyrefly: ignore[missing-import] obj = self._edc_impl.resolve() # Not if the wrapped object do not support tree_map, then it will be @@ -113,7 +113,7 @@ def tree_unflatten( flattened: list[Any], ) -> Any: """`jax.tree_utils` support.""" - import jax # pylint: disable=g-import-not-at-top # pytype: disable=import-error + import jax # pylint: disable=g-import-not-at-top # pyrefly: ignore[missing-import] return jax.tree_util.tree_unflatten(metadata, flattened) diff --git a/etils/edc/frozen_utils_test.py b/etils/edc/frozen_utils_test.py index bc0e1fec..68e0da61 100644 --- a/etils/edc/frozen_utils_test.py +++ b/etils/edc/frozen_utils_test.py @@ -59,15 +59,15 @@ def test_unfrozen_call_twice(): # Can't call frozen on frozen object with pytest.raises(ValueError, match='can only be called after'): - x_origin.frozen() # pytype: disable=attribute-error + x_origin.frozen() # pyrefly: ignore[missing-attribute] # Can call unfrozen twice on the original object - x = x_origin.unfrozen() # pytype: disable=attribute-error - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] # Can't call unfrozen on unfrozen objects with pytest.raises(ValueError, match='Object is already unfrozen'): - x.unfrozen() # pytype: disable=attribute-error + x.unfrozen() y = x.y x.x = 123 @@ -105,7 +105,7 @@ def test_unfrozen_call_twice(): y.x = 123 # Cannot call unfrozen twice on the original object - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] assert x_freezed == A(x=123, y=A(x=456)) assert x_origin == A(y=A(x=456)) @@ -114,7 +114,7 @@ def test_unfrozen_call_twice(): def test_unfrozen_original_obj_non_mutated(): x_origin = A(x='abc', y=A(y='def')) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] assert x.y.y == 'def' x.x = 123 x.y = 456 @@ -131,7 +131,7 @@ def test_unfrozen_original_obj_non_mutated(): def test_unfrozen_assigned_twice(): x_origin = A(x='abc', y=B(y='def')) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.x = x.y # Assign the `B()` to 2 different attributes x.y.x = 123 # Updating one update the other @@ -143,7 +143,7 @@ def test_unfrozen_assigned_twice(): def test_unfrozen_nested(): x_origin = A(x=B(x=A(x=123))) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] val = A() x.x.x.x = A(x=val, y=A(y=123)) x.x.x.x.y.y = 456 @@ -156,7 +156,7 @@ def test_unfrozen_nested(): def test_unfrozen_inheritance(): x_origin = C(x=B(x=A(x=123))) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] val = A() x.x.x.x = A(x=val, y=A(y=123)) x.x.x.x.y.y = 456 @@ -171,7 +171,7 @@ def test_unfrozen_inheritance(): def test_unfrozen_assigned_object(): x_origin = A(x=A(x=123)) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.x = A(x=456) x.x.y = 567 # pyrefly: ignore[read-only] @@ -184,7 +184,7 @@ def test_unfrozen_same_object(): a = A(x=1, y=2) x_origin = A(x=a, y=a) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.x.x = 10 x = x.frozen() @@ -194,7 +194,7 @@ def test_unfrozen_same_object(): assert x == A(x=A(x=10, y=2), y=A(x=1, y=2)) assert x_origin == A(x=A(x=1, y=2), y=A(x=1, y=2)) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.x.x = 10 x.y.y = 20 x = x.frozen() @@ -206,7 +206,7 @@ def test_unfrozen_same_object(): def test_unfrozen_chex(): x_origin = E(x=E(x=A(y=E(x=123)))) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.x.x.y.x = 456 x = x.frozen() assert isinstance(x, E) @@ -221,7 +221,7 @@ def test_unfrozen_chex(): def test_unfrozen_no_updates(): x_origin = A(x=A(x=A(x=123)), y=A(y=A(y=456))) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] assert x.x.x.x == 123 # Access read-only does not trigger `.replace` assert x.y.y.y == 456 x = x.frozen() @@ -234,7 +234,7 @@ def test_unfrozen_no_updates(): assert x_origin.y.y is x.y.y assert x_origin.y is x.y - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() x.x.x.x = 678 assert x.y.y.y == 456 x = x.frozen() @@ -250,7 +250,7 @@ def test_unfrozen_no_updates(): def test_unfrozen_tree_map(): x_origin = E(x=E(x=456)) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] y = jax.tree_util.tree_map(lambda v: v * 10, x.x) assert isinstance(y, E) @@ -268,7 +268,7 @@ def test_unfrozen_tree_map(): @pytest.mark.xfail def test_unfrozen_cicle(): x_origin = A(y=B(y=123)) - x = x_origin.unfrozen() # pytype: disable=attribute-error + x = x_origin.unfrozen() # pyrefly: ignore[missing-attribute] x.y.y = x x = x.frozen() assert x.y.y is x diff --git a/etils/edc/helpers.py b/etils/edc/helpers.py index ae54b866..4cd496ff 100644 --- a/etils/edc/helpers.py +++ b/etils/edc/helpers.py @@ -128,7 +128,7 @@ def _get_type_hints(cls, *, include_extras: bool = False) -> _Hints: lines = [f' * {k}: {v!r}' for k, v in cls.__annotations__.items()] lines = '\n'.join(lines) - epy.reraise(e, prefix=msg + lines + '\n') # pytype: disable=bad-return-type + epy.reraise(e, prefix=msg + lines + '\n') def _get_type_hints_fix(cls, *, include_extras: bool = False) -> _Hints: diff --git a/etils/epath/abstract_path.py b/etils/epath/abstract_path.py index 2c7204c4..0077d396 100644 --- a/etils/epath/abstract_path.py +++ b/etils/epath/abstract_path.py @@ -78,7 +78,7 @@ def __new__(cls, *args: PathLike) -> Self: def format(self, *args: Any, **kwargs: Any) -> Self: """Apply `str.format()` to the path.""" - return type(self)(os.fspath(self).format(*args, **kwargs)) # pytype: disable=not-instantiable + return type(self)(os.fspath(self).format(*args, **kwargs)) # ====== Read-only methods ====== @@ -139,7 +139,7 @@ def walk( def expanduser(self) -> Self: """Returns a new path with expanded `~` and `~user` constructs.""" - if '~' not in self.parts: # pytype: disable=attribute-error + if '~' not in self.parts: return self raise NotImplementedError @@ -235,21 +235,18 @@ def touch(self, mode: int = 0o666, exist_ok: bool = True) -> None: raise FileExistsError(f'{self} already exists.') self.write_text('') - # pytype: disable=bad-return-type @abstractmethod - def rename(self, target: PathLike) -> Self: + def rename(self, target: PathLike) -> Self: # pyrefly: ignore[bad-return] """Renames the path.""" @abstractmethod - def replace(self, target: PathLike) -> Self: + def replace(self, target: PathLike) -> Self: # pyrefly: ignore[bad-return] """Overwrites the destination path.""" @abstractmethod - def copy(self, dst: PathLike, overwrite: bool = False) -> Self: + def copy(self, dst: PathLike, overwrite: bool = False) -> Self: # pyrefly: ignore[bad-return] """Copy the current file to the given destination.""" - # pytype: enable=bad-return-type - # ====== Other ====== @classmethod @@ -265,7 +262,7 @@ def __get_pydantic_core_schema__(cls, source_type, handler): Returns: pydantic_core.CoreSchema """ - from pydantic_core import core_schema # pylint: disable=g-import-not-at-top # pytype: disable=import-error + from pydantic_core import core_schema # pylint: disable=g-import-not-at-top # pyrefly: ignore[missing-import] del source_type # Unused in this implementation. del handler # Unused in this implementation. diff --git a/etils/epath/backend.py b/etils/epath/backend.py index 87f38980..adcb95ff 100644 --- a/etils/epath/backend.py +++ b/etils/epath/backend.py @@ -251,7 +251,7 @@ class _TfBackend(Backend): @property def tf(self): try: - import tensorflow # pylint: disable=g-import-not-at-top # pytype: disable=import-error + import tensorflow # pylint: disable=g-import-not-at-top # pyrefly: ignore[missing-source-for-stubs] except ImportError as e: raise ImportError( f'{e}. To use epath.Path with gs://, TensorFlow should be installed.' @@ -268,7 +268,7 @@ def open( # pyrefly: ignore[bad-override] path: PathLike, mode: str, ) -> Iterator[typing.IO[Union[str, bytes]]]: - with self.gfile.GFile(path, mode) as f: # pytype: disable=bad-return-type + with self.gfile.GFile(path, mode) as f: try: yield f except self.tf.errors.NotFoundError as e: @@ -284,7 +284,7 @@ def listdir(self, path: PathLike) -> list[str]: return self.gfile.listdir(path) def glob(self, path: PathLike) -> list[str]: - return self.gfile.glob(path) # pyrefly: ignore[bad-argument-type] + return self.gfile.glob(path) def walk( self, @@ -430,6 +430,8 @@ def _get_filesystem(self, name: str) -> fsspec.AbstractFileSystem: raise ImportError( "To use epath.Path with gs://, fsspec should be installed.'" ) from e + if name in ('gcs', 'gs'): + return fsspec.filesystem(name, finalize_on_close=True) return fsspec.filesystem(name) def fs(self, path: PathLike) -> fsspec.AbstractFileSystem: @@ -470,7 +472,7 @@ def walk( top_down: bool = True, on_error: Callable[[OSError], object] | None = None, ) -> Iterator[tuple[PathLike, list[str], list[str]]]: - yield from self.fs(top).walk( # pytype: disable=bad-return-type + yield from self.fs(top).walk( top, topdown=top_down, # default behavior for pathlib.Path.walk @@ -506,6 +508,10 @@ def mkdir( 'mkdir with custom `mode=` not supported for FileSystemSpec backend.' ' Please open an issue.' ) + if not exist_ok and self.exists(path): + raise FileExistsError( + f'The operation failed because the specified {path=} already exists.' + ) try: return self.fs(path).mkdir(path, create_parents=False) except FileExistsError: diff --git a/etils/epath/backend_test.py b/etils/epath/backend_test.py index cb4d4230..60d329f7 100644 --- a/etils/epath/backend_test.py +++ b/etils/epath/backend_test.py @@ -22,6 +22,7 @@ import pathlib import pwd from typing import Dict, Union +from unittest import mock from etils import epath from etils import epy @@ -322,6 +323,14 @@ def _test_mkdir_exists_not_ok( with pytest.raises(FileNotFoundError, match='No such file or directory'): backend.mkdir(tmp_path / 'nested/non-existing', exist_ok=False) + if backend == epath.backend.fsspec_backend: + # Cloud fsspec backends (e.g. gcsfs, s3fs) implement `mkdir` as a no-op. + fs = epath.backend.fsspec_backend.fs(tmp_path) + with mock.patch.object(fs, 'mkdir', return_value=None): + backend.mkdir(tmp_path / 'folder', exist_ok=True) + with pytest.raises(FileExistsError): + backend.mkdir(tmp_path / 'folder', exist_ok=False) + def _test_rmtree( backend: epath.backend.Backend, diff --git a/etils/epath/flags.py b/etils/epath/flags.py index dcf06cbf..c87cb090 100644 --- a/etils/epath/flags.py +++ b/etils/epath/flags.py @@ -113,4 +113,4 @@ def serialize(self, value): _PathSerializer(), required=required, **kwargs, - ) # pytype: disable=bad-return-type + ) diff --git a/etils/epath/gpath.py b/etils/epath/gpath.py index b175d186..9e6cebea 100644 --- a/etils/epath/gpath.py +++ b/etils/epath/gpath.py @@ -55,6 +55,31 @@ def _epath_use_tf() -> bool: '0', ] + +@functools.lru_cache(maxsize=1) +def _get_storage_control_client() -> Any: + """Returns a cached `StorageControlClient` instance.""" + # pylint: disable=g-import-not-at-top + # pyrefly: ignore[missing-import] + from google.cloud import storage_control_v2 + + return storage_control_v2.StorageControlClient() + + +@functools.lru_cache(maxsize=128) +def _is_zonal_gcs_bucket(bucket: str) -> bool: + """Returns True if `bucket` is a GCS Zonal Rapid bucket.""" + if not bucket: + return False + try: + layout = _get_storage_control_client().get_storage_layout( + name=f'projects/_/buckets/{bucket}/storageLayout' + ) + return layout.location_type == 'zone' + except Exception: # pylint: disable=broad-except + # Missing SDK, credentials, or permissions fall back to default backend. + return False + _PREFIX_TO_BACKEND = { 'gs': backend_lib.fsspec_backend, 's3': backend_lib.fsspec_backend, @@ -116,6 +141,9 @@ def _backend(self) -> backend_lib.Backend: # for retro-compatibility, because needed dependencies (gcsfs or s3fs) # may not be installed. fsspec_backend was indeed introduced later. if _is_tf_installed() and self._uri_scheme is not None: + bucket = self.parts[2] if len(self.parts) > 2 else '' + if self._uri_scheme == 'gs' and _is_zonal_gcs_bucket(bucket): + return backend_lib.fsspec_backend return backend_lib.tf_backend return backend except KeyError: @@ -232,7 +260,7 @@ def unlink(self, missing_ok: bool = False) -> None: if not missing_ok: raise - def open( # pytype: disable=signature-mismatch # overriding-parameter-count-checks + def open( # pyrefly: ignore[bad-override] self, mode: str = 'r', *, diff --git a/etils/epath/gpath_test.py b/etils/epath/gpath_test.py index 03887c4d..86348c43 100644 --- a/etils/epath/gpath_test.py +++ b/etils/epath/gpath_test.py @@ -18,6 +18,7 @@ import os import pathlib +import types from unittest import mock from etils import epath @@ -437,17 +438,27 @@ def test_use_backend(): gs_backend = epath.backend.tf_backend # pylint: disable=unused-variable loc_backend = epath.backend.os_backend # pylint: disable=unused-variable with mock.patch.object(epath.gpath, '_is_tf_installed', return_value=True): - assert epath.gpath._get_backend(gs_path, gs_path) == gs_backend + with mock.patch.object( + epath.gpath, '_is_zonal_gcs_bucket', return_value=False + ): + assert epath.gpath._get_backend(gs_path, gs_path) == gs_backend + with mock.patch.object( + epath.gpath, '_is_zonal_gcs_bucket', return_value=True + ): + assert ( + epath.gpath._get_backend(gs_path, gs_path) + == epath.backend.fsspec_backend + ) with mock.patch.object(epath.gpath, '_is_tf_installed', return_value=False): assert ( epath.gpath._get_backend(gs_path, gs_path) == epath.backend.fsspec_backend ) - assert epath.gpath._get_backend(gs_path, gs_path) == gs_backend # pytype: disable=wrong-arg-types - assert epath.gpath._get_backend(gs_path, loc_path) == gs_backend # pytype: disable=wrong-arg-types - assert epath.gpath._get_backend(loc_path, gs_path) == gs_backend # pytype: disable=wrong-arg-types - assert epath.gpath._get_backend(loc_path, loc_path) == loc_backend # pytype: disable=wrong-arg-types + assert epath.gpath._get_backend(gs_path, gs_path) == gs_backend # pyrefly: ignore[bad-argument-type] + assert epath.gpath._get_backend(gs_path, loc_path) == gs_backend # pyrefly: ignore[bad-argument-type] + assert epath.gpath._get_backend(loc_path, gs_path) == gs_backend # pyrefly: ignore[bad-argument-type] + assert epath.gpath._get_backend(loc_path, loc_path) == loc_backend # pyrefly: ignore[bad-argument-type] @epy.testing.non_hermetic @@ -463,3 +474,40 @@ def test_relative_to(): assert path.relative_to('gs://bucket/dir') == epath.Path('subdir') with pytest.raises(ValueError, match='not in the subpath'): path.relative_to('gs://bucket/other-dir') + + +@pytest.mark.parametrize( + 'bucket,layout,expected', + [ + ('zonal-bucket', types.SimpleNamespace(location_type='zone'), True), + ( + 'regional-bucket', + types.SimpleNamespace(location_type='region'), + False, + ), + ('error-bucket', RuntimeError('offline'), False), + ('', None, False), + ], +) +def test_is_zonal_gcs_bucket( + bucket: str, + layout: object, + expected: bool, + monkeypatch: pytest.MonkeyPatch, +): + epath.gpath._is_zonal_gcs_bucket.cache_clear() + + def _get_layout(name: str): + assert bucket and name == f'projects/_/buckets/{bucket}/storageLayout' + if isinstance(layout, Exception): + raise layout + return layout + + client = types.SimpleNamespace(get_storage_layout=_get_layout) + monkeypatch.setattr( + epath.gpath, '_get_storage_control_client', lambda: client + ) + try: + assert epath.gpath._is_zonal_gcs_bucket(bucket) is expected + finally: + epath.gpath._is_zonal_gcs_bucket.cache_clear() diff --git a/etils/epath/register.py b/etils/epath/register.py index f96c7fbb..1874d5f8 100644 --- a/etils/epath/register.py +++ b/etils/epath/register.py @@ -71,9 +71,9 @@ def register_path_cls_decorator(cls: _T) -> _T: elif isinstance(path_cls_or_uri_prefix, (list, tuple)): for uri_prefix in path_cls_or_uri_prefix: _URI_PREFIXES_TO_CLS[uri_prefix] = cls - return register_path_cls(cls) # pytype: disable=bad-return-type + return register_path_cls(cls) - return register_path_cls_decorator # pytype: disable=bad-return-type + return register_path_cls_decorator else: _PATHLIKE_CLS = _PATHLIKE_CLS + (path_cls_or_uri_prefix,) return path_cls_or_uri_prefix @@ -97,13 +97,13 @@ def make_path(path: PathLike) -> abstract_path.Path: uri_splits = path.split('://', maxsplit=1) if len(uri_splits) > 1: # str is URI (e.g. `gs://`, `github://`,...) # On windows, `PosixGPath` is created for `gs://` paths - return _URI_PREFIXES_TO_CLS[uri_splits[0] + '://'](path) # pytype: disable=bad-return-type + return _URI_PREFIXES_TO_CLS[uri_splits[0] + '://'](path) elif is_windows: return gpath.WindowsGPath(path) else: return gpath.PosixGPath(path) elif isinstance(path, _PATHLIKE_CLS): - return path # Forward resource path, gpath,... as-is # pytype: disable=bad-return-type + return path # Forward resource path, gpath,... as-is elif isinstance(path, os.PathLike): # Other `os.fspath` compatible objects path_cls = gpath.WindowsGPath if is_windows else gpath.PosixGPath return path_cls(path) diff --git a/etils/epath/resource_utils.py b/etils/epath/resource_utils.py index c82488da..79d0d1b1 100644 --- a/etils/epath/resource_utils.py +++ b/etils/epath/resource_utils.py @@ -57,29 +57,29 @@ def __fspath__(self) -> str: def __eq__(self, other) -> bool: return ( type(self) == type(other) # pylint: disable=unidiomatic-typecheck - and self.root == other.root # pytype: disable=attribute-error - and self.at == other.at # pytype: disable=attribute-error + and self.root == other.root + and self.at == other.at ) def __hash__(self) -> int: - return hash((self.root, self.at)) # pytype: disable=attribute-error + return hash((self.root, self.at)) if sys.version_info < (3, 10): # Required due to: https://bugs.python.org/issue42043 def _next(self, at) -> 'ResourcePath': # pylint: disable=g-wrong-blank-lines - return type(self)(self.root, at) # pytype: disable=attribute-error + return type(self)(self.root, at) # Before 3.10, joinpath only accept a single arg def joinpath(self, *other): """Overwrite `joinpath` to be consistent with `pathlib.Path`.""" - next_ = posixpath.join(self.at, *other) # pytype: disable=attribute-error - return self._next(self.root.resolve_dir(next_)) # pytype: disable=attribute-error + next_ = posixpath.join(self.at, *other) + return self._next(self.root.resolve_dir(next_)) if sys.version_info < (3, 11): @property def suffix(self): - return pathlib.Path(self.at).suffix or self.filename.suffix # pytype: disable=attribute-error + return pathlib.Path(self.at).suffix or self.filename.suffix def resource_path(package: Union[str, types.ModuleType]) -> abstract_path.Path: @@ -110,7 +110,7 @@ def resource_path(package: Union[str, types.ModuleType]) -> abstract_path.Path: The read-only path to the root module directory """ try: - path = importlib_resources.files(package) # pytype: disable=module-attr + path = importlib_resources.files(package) except AttributeError: path = None is_adhoc = True diff --git a/etils/epath/testing.py b/etils/epath/testing.py index e1eb3303..c421d99b 100644 --- a/etils/epath/testing.py +++ b/etils/epath/testing.py @@ -108,7 +108,7 @@ def rename(self, path: PathLike, dst: PathLike) -> None: def replace(self, path: PathLike, dst: PathLike) -> None: return self._get_fn('replace')(path, dst) - def copy(self, path: PathLike, dst: PathLike, *, overwrite: bool) -> None: # pytype: disable=signature-mismatch # overriding-parameter-count-checks + def copy(self, path: PathLike, dst: PathLike, *, overwrite: bool) -> None: # pyrefly: ignore[bad-override] return self._get_fn('copy')(path, dst, overwrite=overwrite) def stat(self, path: PathLike) -> stat_utils.StatResult: diff --git a/pyproject.toml b/pyproject.toml index dc7fc6f5..3afd4c25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,7 @@ epath = [ "etils[epy]", ] epath-gcs = [ - "gcsfs", + "gcsfs>=2026.4.0", "etils[epath]", ] epath-s3 = [