Skip to content

Commit 5930d11

Browse files
headtr1ckheadtr1ck
authored
Buffer types (#9787)
* add ReadBuffer Protocol for open_mfdataset * finally fix LSP violation * move import out of TYPE_CHECKING
1 parent d5f84dd commit 5930d11

File tree

11 files changed

+123
-67
lines changed

11 files changed

+123
-67
lines changed

‎xarray/backends/api.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -54,14 +54,14 @@
5454
from dask.delayed import Delayed
5555
except ImportError:
5656
Delayed = None # type: ignore[assignment, misc]
57-
from io import BufferedIOBase
5857

5958
from xarray.backends.common import BackendEntrypoint
6059
from xarray.core.types import (
6160
CombineAttrsOptions,
6261
CompatOptions,
6362
JoinOptions,
6463
NestedSequence,
64+
ReadBuffer,
6565
T_Chunks,
6666
)
6767

@@ -474,7 +474,7 @@ def _datatree_from_backend_datatree(
474474

475475

476476
def open_dataset(
477-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
477+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
478478
*,
479479
engine: T_Engine = None,
480480
chunks: T_Chunks = None,
@@ -691,7 +691,7 @@ def open_dataset(
691691

692692

693693
def open_dataarray(
694-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
694+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
695695
*,
696696
engine: T_Engine | None = None,
697697
chunks: T_Chunks | None = None,
@@ -896,7 +896,7 @@ def open_dataarray(
896896

897897

898898
def open_datatree(
899-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
899+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
900900
*,
901901
engine: T_Engine = None,
902902
chunks: T_Chunks = None,
@@ -1111,7 +1111,7 @@ def open_datatree(
11111111

11121112

11131113
def open_groups(
1114-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
1114+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
11151115
*,
11161116
engine: T_Engine = None,
11171117
chunks: T_Chunks = None,
@@ -1137,10 +1137,6 @@ def open_groups(
11371137
and cannot be opened directly with ``open_datatree``. It is encouraged to use this function to inspect your data,
11381138
then make the necessary changes to make the structure coercible to a `DataTree` object before calling `DataTree.from_dict()` and proceeding with your analysis.
11391139
1140-
Parameters
1141-
----------
1142-
filename_or_obj : str, Path, file-like, or DataStore
1143-
Strings and Path objects are interpreted as a path to a netCDF file.
11441140
Parameters
11451141
----------
11461142
filename_or_obj : str, Path, file-like, or DataStore
@@ -1338,7 +1334,10 @@ def open_groups(
13381334

13391335

13401336
def open_mfdataset(
1341-
paths: str | os.PathLike | NestedSequence[str | os.PathLike],
1337+
paths: str
1338+
| os.PathLike
1339+
| ReadBuffer
1340+
| NestedSequence[str | os.PathLike | ReadBuffer],
13421341
chunks: T_Chunks | None = None,
13431342
concat_dim: (
13441343
str
@@ -1541,7 +1540,7 @@ def open_mfdataset(
15411540
if not paths:
15421541
raise OSError("no files to open")
15431542

1544-
paths1d: list[str]
1543+
paths1d: list[str | ReadBuffer]
15451544
if combine == "nested":
15461545
if isinstance(concat_dim, str | DataArray) or concat_dim is None:
15471546
concat_dim = [concat_dim] # type: ignore[assignment]

‎xarray/backends/common.py

Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,19 @@
66
import traceback
77
from collections.abc import Iterable, Mapping, Sequence
88
from glob import glob
9-
from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast, overload
9+
from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, overload
1010

1111
import numpy as np
1212

1313
from xarray.conventions import cf_encoder
1414
from xarray.core import indexing
1515
from xarray.core.datatree import DataTree
16+
from xarray.core.types import ReadBuffer
1617
from xarray.core.utils import FrozenDict, NdimSizeLenMixin, is_remote_uri
1718
from xarray.namedarray.parallelcompat import get_chunked_array_type
1819
from xarray.namedarray.pycompat import is_chunked_array
1920

2021
if TYPE_CHECKING:
21-
from io import BufferedIOBase
22-
2322
from xarray.core.dataset import Dataset
2423
from xarray.core.types import NestedSequence
2524

@@ -65,24 +64,52 @@ def _normalize_path(path: str | os.PathLike | T) -> str | T:
6564
if isinstance(path, str) and not is_remote_uri(path):
6665
path = os.path.abspath(os.path.expanduser(path))
6766

68-
return cast(str, path)
67+
return path # type:ignore [return-value]
6968

7069

7170
@overload
7271
def _find_absolute_paths(
73-
paths: str | os.PathLike | Sequence[str | os.PathLike], **kwargs
72+
paths: str | os.PathLike | Sequence[str | os.PathLike],
73+
**kwargs,
7474
) -> list[str]: ...
7575

7676

77+
@overload
78+
def _find_absolute_paths(
79+
paths: ReadBuffer | Sequence[ReadBuffer],
80+
**kwargs,
81+
) -> list[ReadBuffer]: ...
82+
83+
7784
@overload
7885
def _find_absolute_paths(
7986
paths: NestedSequence[str | os.PathLike], **kwargs
8087
) -> NestedSequence[str]: ...
8188

8289

90+
@overload
8391
def _find_absolute_paths(
84-
paths: str | os.PathLike | NestedSequence[str | os.PathLike], **kwargs
85-
) -> NestedSequence[str]:
92+
paths: NestedSequence[ReadBuffer], **kwargs
93+
) -> NestedSequence[ReadBuffer]: ...
94+
95+
96+
@overload
97+
def _find_absolute_paths(
98+
paths: str
99+
| os.PathLike
100+
| ReadBuffer
101+
| NestedSequence[str | os.PathLike | ReadBuffer],
102+
**kwargs,
103+
) -> NestedSequence[str | ReadBuffer]: ...
104+
105+
106+
def _find_absolute_paths(
107+
paths: str
108+
| os.PathLike
109+
| ReadBuffer
110+
| NestedSequence[str | os.PathLike | ReadBuffer],
111+
**kwargs,
112+
) -> NestedSequence[str | ReadBuffer]:
86113
"""
87114
Find absolute paths from the pattern.
88115
@@ -132,10 +159,12 @@ def _find_absolute_paths(
132159
return sorted(glob(_normalize_path(paths)))
133160
elif isinstance(paths, os.PathLike):
134161
return [_normalize_path(paths)]
162+
elif isinstance(paths, ReadBuffer):
163+
return [paths]
135164

136165
def _normalize_path_list(
137-
lpaths: NestedSequence[str | os.PathLike],
138-
) -> NestedSequence[str]:
166+
lpaths: NestedSequence[str | os.PathLike | ReadBuffer],
167+
) -> NestedSequence[str | ReadBuffer]:
139168
paths = []
140169
for p in lpaths:
141170
if isinstance(p, str | os.PathLike):
@@ -546,10 +575,9 @@ def __repr__(self) -> str:
546575

547576
def open_dataset(
548577
self,
549-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
578+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
550579
*,
551580
drop_variables: str | Iterable[str] | None = None,
552-
**kwargs: Any,
553581
) -> Dataset:
554582
"""
555583
Backend open_dataset method used by Xarray in :py:func:`~xarray.open_dataset`.
@@ -559,7 +587,7 @@ def open_dataset(
559587

560588
def guess_can_open(
561589
self,
562-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
590+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
563591
) -> bool:
564592
"""
565593
Backend open_dataset method used by Xarray in :py:func:`~xarray.open_dataset`.
@@ -569,8 +597,9 @@ def guess_can_open(
569597

570598
def open_datatree(
571599
self,
572-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
573-
**kwargs: Any,
600+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
601+
*,
602+
drop_variables: str | Iterable[str] | None = None,
574603
) -> DataTree:
575604
"""
576605
Backend open_datatree method used by Xarray in :py:func:`~xarray.open_datatree`.
@@ -580,8 +609,9 @@ def open_datatree(
580609

581610
def open_groups_as_dict(
582611
self,
583-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
584-
**kwargs: Any,
612+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
613+
*,
614+
drop_variables: str | Iterable[str] | None = None,
585615
) -> dict[str, Dataset]:
586616
"""
587617
Opens a dictionary mapping from group names to Datasets.

‎xarray/backends/h5netcdf_.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,10 @@
3939
from xarray.core.variable import Variable
4040

4141
if TYPE_CHECKING:
42-
from io import BufferedIOBase
43-
4442
from xarray.backends.common import AbstractDataStore
4543
from xarray.core.dataset import Dataset
4644
from xarray.core.datatree import DataTree
45+
from xarray.core.types import ReadBuffer
4746

4847

4948
class H5NetCDFArrayWrapper(BaseNetCDF4Array):
@@ -395,7 +394,7 @@ class H5netcdfBackendEntrypoint(BackendEntrypoint):
395394

396395
def guess_can_open(
397396
self,
398-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
397+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
399398
) -> bool:
400399
magic_number = try_read_magic_number_from_file_or_path(filename_or_obj)
401400
if magic_number is not None:
@@ -407,9 +406,9 @@ def guess_can_open(
407406

408407
return False
409408

410-
def open_dataset( # type: ignore[override] # allow LSP violation, not supporting **kwargs
409+
def open_dataset(
411410
self,
412-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
411+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
413412
*,
414413
mask_and_scale=True,
415414
decode_times=True,
@@ -456,7 +455,7 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti
456455

457456
def open_datatree(
458457
self,
459-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
458+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
460459
*,
461460
mask_and_scale=True,
462461
decode_times=True,
@@ -499,7 +498,7 @@ def open_datatree(
499498

500499
def open_groups_as_dict(
501500
self,
502-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
501+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
503502
*,
504503
mask_and_scale=True,
505504
decode_times=True,

‎xarray/backends/netCDF4_.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,13 @@
4141
from xarray.core.variable import Variable
4242

4343
if TYPE_CHECKING:
44-
from io import BufferedIOBase
45-
4644
from h5netcdf.core import EnumType as h5EnumType
4745
from netCDF4 import EnumType as ncEnumType
4846

4947
from xarray.backends.common import AbstractDataStore
5048
from xarray.core.dataset import Dataset
5149
from xarray.core.datatree import DataTree
50+
from xarray.core.types import ReadBuffer
5251

5352
# This lookup table maps from dtype.byteorder to a readable endian
5453
# string used by netCDF4.
@@ -627,7 +626,7 @@ class NetCDF4BackendEntrypoint(BackendEntrypoint):
627626

628627
def guess_can_open(
629628
self,
630-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
629+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
631630
) -> bool:
632631
if isinstance(filename_or_obj, str) and is_remote_uri(filename_or_obj):
633632
return True
@@ -642,9 +641,9 @@ def guess_can_open(
642641

643642
return False
644643

645-
def open_dataset( # type: ignore[override] # allow LSP violation, not supporting **kwargs
644+
def open_dataset(
646645
self,
647-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
646+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
648647
*,
649648
mask_and_scale=True,
650649
decode_times=True,
@@ -693,7 +692,7 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti
693692

694693
def open_datatree(
695694
self,
696-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
695+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
697696
*,
698697
mask_and_scale=True,
699698
decode_times=True,
@@ -735,7 +734,7 @@ def open_datatree(
735734

736735
def open_groups_as_dict(
737736
self,
738-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
737+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
739738
*,
740739
mask_and_scale=True,
741740
decode_times=True,

‎xarray/backends/plugins.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,9 @@
1414
if TYPE_CHECKING:
1515
import os
1616
from importlib.metadata import EntryPoint, EntryPoints
17-
from io import BufferedIOBase
1817

1918
from xarray.backends.common import AbstractDataStore
19+
from xarray.core.types import ReadBuffer
2020

2121
STANDARD_BACKENDS_ORDER = ["netcdf4", "h5netcdf", "scipy"]
2222

@@ -138,7 +138,7 @@ def refresh_engines() -> None:
138138

139139

140140
def guess_engine(
141-
store_spec: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
141+
store_spec: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
142142
) -> str | type[BackendEntrypoint]:
143143
engines = list_engines()
144144

‎xarray/backends/pydap_.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,9 @@
2626

2727
if TYPE_CHECKING:
2828
import os
29-
from io import BufferedIOBase
3029

3130
from xarray.core.dataset import Dataset
31+
from xarray.core.types import ReadBuffer
3232

3333

3434
class PydapArrayWrapper(BackendArray):
@@ -166,13 +166,13 @@ class PydapBackendEntrypoint(BackendEntrypoint):
166166

167167
def guess_can_open(
168168
self,
169-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
169+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
170170
) -> bool:
171171
return isinstance(filename_or_obj, str) and is_remote_uri(filename_or_obj)
172172

173-
def open_dataset( # type: ignore[override] # allow LSP violation, not supporting **kwargs
173+
def open_dataset(
174174
self,
175-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
175+
filename_or_obj: str | os.PathLike[Any] | ReadBuffer | AbstractDataStore,
176176
*,
177177
mask_and_scale=True,
178178
decode_times=True,

0 commit comments

Comments
 (0)