Support LazyModule setattr and delattr

PiperOrigin-RevId: 737044617

Author: The etils Authors

CHANGELOG.md

@@ -8,8 +8,10 @@ Changelog follow https://keepachangelog.com/ format.
 
 ## [Unreleased]
 
-*  `ecolab.adhoc`:
+*   `ecolab.adhoc`:
     *   Better error message for adhoc invalidate with `epy.reraise`
+*   `epy`:
+    *   Support `__setattr__` and `__delattr__` on LazyModules.
 
 ## [1.12.2] - 2025-03-10
 

etils/epy/lazy_imports_utils.py

@@ -29,6 +29,7 @@
 import dataclasses
 import functools
 import importlib
+import inspect
 import sys
 import threading
 import types
@@ -40,6 +41,8 @@
 
 _ErrorCallback = Callable[[Exception], None]
 _SuccessCallback = Callable[[str], None]
+# Eagerly resolve this import since it's needed in setattr.
+_getattr_static = inspect.getattr_static
 
 # Store a lock per module to avoid problems with multiple threads trying to
 # import the same module at the same time.
@@ -55,6 +58,7 @@ class LazyModule:
   error_callback: str | _ErrorCallback | None
   success_callback: _SuccessCallback | None
   _submodules: dict[str, LazyModule] = dataclasses.field(default_factory=dict)
+  _initialized: bool = dataclasses.field(default=False, init=False)
 
   def __post_init__(self):
     if self.adhoc_kwargs is not None:
@@ -66,6 +70,7 @@ def __post_init__(self):
       self.adhoc_kwargs.pop("reload_workspace", None)
       self.adhoc_kwargs.pop("cell_autoreload", None)
       self.adhoc_kwargs.pop("restrict_reload", None)
+    self._initialized = True
 
   @functools.cached_property
   def _module(self) -> types.ModuleType:
@@ -117,7 +122,20 @@ def __getattr__(self, name: str) -> Any:
     else:
       return getattr(self._module, name)
 
-  # TODO(epot): Also support __setattr__
+  def __setattr__(self, name: str, value: Any) -> None:
+    # Avoid regular `__getattr__` path during initialization.
+    if _getattr_static(self, "_initialized"):
+      # Trigger import first to overwrite the old attribute if it exists.
+      setattr(self._module, name, value)
+    else:
+      super().__setattr__(name, value)
+
+  def __delattr__(self, name: str):
+    if name in self._submodules:
+      del self._submodules[name]
+    # Always delete from underlying module so that lazy import doesn't replace
+    # a deleted attribute.
+    delattr(self._module, name)
 
 
 def _register_submodule(module: LazyModule, name: str) -> LazyModule:

etils/epy/lazy_imports_utils_test.py

@@ -64,6 +64,37 @@ def test_import_with_alias():
   assert 'tensorflow_datasets' in sys.modules
 
 
+def test_setattr():
+  assert 'tensorflow_datasets' not in sys.modules
+  with epy.lazy_imports():
+    import tensorflow_datasets as tfds  # pylint: disable=g-import-not-at-top
+  assert 'tensorflow_datasets' not in sys.modules
+  tfds.features = 'foo'
+  assert 'tensorflow_datasets' in sys.modules
+  assert tfds.features == 'foo'
+
+
+def test_delattr():
+  assert 'tensorflow_datasets' not in sys.modules
+  with epy.lazy_imports():
+    import tensorflow_datasets as tfds  # pylint: disable=g-import-not-at-top
+  assert 'tensorflow_datasets' not in sys.modules
+  del tfds.features
+  assert 'tensorflow_datasets' in sys.modules
+  with pytest.raises(AttributeError):
+    _ = tfds.features
+
+
+def test_delattr_submodule():
+  assert 'tensorflow_datasets' not in sys.modules
+  with epy.lazy_imports():
+    import tensorflow_datasets.core  # pylint: disable=g-import-not-at-top
+  assert 'tensorflow_datasets' not in sys.modules
+  del tensorflow_datasets.core
+  with pytest.raises(AttributeError):
+    _ = tensorflow_datasets.core
+
+
 def test_error_callback():
   success_callback = mock.MagicMock()
   error_callback = mock.MagicMock()