Skip to content

Commit

Permalink
Make list_chunkmanagers more resilient to broken entrypoints
Browse files Browse the repository at this point in the history
As I'm a developing my custom chunk manager, I'm often checking out
between my development branch and production branch breaking the
entrypoint.

This made xarray impossible to import unless I re-ran `pip install -e . -vv`
which is somewhat tiring.

This should help xarray be more resilient in other software's bugs in
case they install malformed entrypoints

Example:

```python
>>> from xarray.core.parallelcompat import list_chunkmanagers

>>> list_chunkmanagers()
<ipython-input-3-19326f4950bc>:1: UserWarning: Failed to load entrypoint MyChunkManager due to No module named 'my.array._chunkmanager'. Skipping.
  list_chunkmanagers()
{'dask': <xarray.core.daskmanager.DaskManager at 0x7f5b826231c0>}
```
  • Loading branch information
hmaarrfk committed Feb 13, 2024
1 parent 0ad7fa7 commit bfb7cdd
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 3 deletions.
13 changes: 10 additions & 3 deletions xarray/core/parallelcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import numpy as np

from xarray.core.pycompat import is_chunked_array
from xarray.core.utils import emit_user_level_warning

T_ChunkedArray = TypeVar("T_ChunkedArray", bound=Any)

Expand Down Expand Up @@ -57,9 +58,15 @@ def load_chunkmanagers(
) -> dict[str, ChunkManagerEntrypoint]:
"""Load entrypoints and instantiate chunkmanagers only once."""

loaded_entrypoints = {
entrypoint.name: entrypoint.load() for entrypoint in entrypoints
}
loaded_entrypoints = {}
for entrypoint in entrypoints:
try:
loaded_entrypoints[entrypoint.name] = entrypoint.load()
except ModuleNotFoundError as e:
emit_user_level_warning(
f"Failed to load chunk manager entrypoint {entrypoint.name} due to {e}. Skipping.",
)
pass

available_chunkmanagers = {
name: chunkmanager()
Expand Down
9 changes: 9 additions & 0 deletions xarray/tests/test_parallelcompat.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

from importlib.metadata import EntryPoint
from typing import Any

import numpy as np
Expand All @@ -11,6 +12,7 @@
get_chunked_array_type,
guess_chunkmanager,
list_chunkmanagers,
load_chunkmanagers,
)
from xarray.core.types import T_Chunks, T_DuckArray, T_NormalizedChunks
from xarray.tests import has_dask, requires_dask
Expand Down Expand Up @@ -217,3 +219,10 @@ def test_raise_on_mixed_array_types(self, register_dummy_chunkmanager) -> None:

with pytest.raises(TypeError, match="received multiple types"):
get_chunked_array_type(*[dask_arr, dummy_arr])

def test_bogus_entrypoint():
# Create a bogus entry-point as if the user broke their setup.cfg
# or is actively developing their new chunk manager
entry_point = EntryPoint("bogus", "xarray.bogus.doesnotwork", "xarray.chunkmanagers")
with pytest.warns(UserWarning, match="Failed to load chunk manager"):
assert len(load_chunkmanagers([entry_point])) == 0

0 comments on commit bfb7cdd

Please sign in to comment.