12
12
from zarr .hierarchy import group as _create_group
13
13
from zarr .hierarchy import open_group
14
14
from zarr .meta import json_dumps , json_loads
15
- from zarr .storage import contains_array , contains_group
15
+ from zarr .storage import contains_array , contains_group , Store
16
16
from zarr .util import TreeViewer , buffer_size , normalize_storage_path
17
17
18
+ from typing import Union
19
+
20
+ StoreLike = Union [Store , str , None ]
21
+
18
22
19
23
# noinspection PyShadowingBuiltins
20
- def open (store = None , mode = 'a' , ** kwargs ):
24
+ def open (store : StoreLike = None , mode : str = "a" , ** kwargs ):
21
25
"""Convenience function to open a group or array using file-mode-like semantics.
22
26
23
27
Parameters
24
28
----------
25
- store : MutableMapping or string, optional
29
+ store : Store or string, optional
26
30
Store or path to directory in file system or name of zip file.
27
31
mode : {'r', 'r+', 'a', 'w', 'w-'}, optional
28
32
Persistence mode: 'r' means read only (must exist); 'r+' means
@@ -75,32 +79,33 @@ def open(store=None, mode='a', **kwargs):
75
79
clobber = mode == 'w'
76
80
# we pass storage options explicitly, since normalize_store_arg might construct
77
81
# a store if the input is a fsspec-compatible URL
78
- store = normalize_store_arg (store , clobber = clobber ,
79
- storage_options = kwargs .pop ("storage_options" , {}))
82
+ _store : Store = normalize_store_arg (
83
+ store , clobber = clobber , storage_options = kwargs .pop ("storage_options" , {})
84
+ )
80
85
path = normalize_storage_path (path )
81
86
82
87
if mode in {'w' , 'w-' , 'x' }:
83
88
if 'shape' in kwargs :
84
- return open_array (store , mode = mode , ** kwargs )
89
+ return open_array (_store , mode = mode , ** kwargs )
85
90
else :
86
- return open_group (store , mode = mode , ** kwargs )
91
+ return open_group (_store , mode = mode , ** kwargs )
87
92
88
93
elif mode == "a" :
89
- if "shape" in kwargs or contains_array (store , path ):
90
- return open_array (store , mode = mode , ** kwargs )
94
+ if "shape" in kwargs or contains_array (_store , path ):
95
+ return open_array (_store , mode = mode , ** kwargs )
91
96
else :
92
- return open_group (store , mode = mode , ** kwargs )
97
+ return open_group (_store , mode = mode , ** kwargs )
93
98
94
99
else :
95
- if contains_array (store , path ):
96
- return open_array (store , mode = mode , ** kwargs )
97
- elif contains_group (store , path ):
98
- return open_group (store , mode = mode , ** kwargs )
100
+ if contains_array (_store , path ):
101
+ return open_array (_store , mode = mode , ** kwargs )
102
+ elif contains_group (_store , path ):
103
+ return open_group (_store , mode = mode , ** kwargs )
99
104
else :
100
105
raise PathNotFoundError (path )
101
106
102
107
103
- def save_array (store , arr , ** kwargs ):
108
+ def save_array (store : StoreLike , arr , ** kwargs ):
104
109
"""Convenience function to save a NumPy array to the local file system, following a
105
110
similar API to the NumPy save() function.
106
111
@@ -132,16 +137,16 @@ def save_array(store, arr, **kwargs):
132
137
133
138
"""
134
139
may_need_closing = isinstance (store , str )
135
- store = normalize_store_arg (store , clobber = True )
140
+ _store : Store = normalize_store_arg (store , clobber = True )
136
141
try :
137
- _create_array (arr , store = store , overwrite = True , ** kwargs )
142
+ _create_array (arr , store = _store , overwrite = True , ** kwargs )
138
143
finally :
139
- if may_need_closing and hasattr ( store , 'close' ) :
144
+ if may_need_closing :
140
145
# needed to ensure zip file records are written
141
- store .close ()
146
+ _store .close ()
142
147
143
148
144
- def save_group (store , * args , ** kwargs ):
149
+ def save_group (store : StoreLike , * args , ** kwargs ):
145
150
"""Convenience function to save several NumPy arrays to the local file system, following a
146
151
similar API to the NumPy savez()/savez_compressed() functions.
147
152
@@ -203,21 +208,21 @@ def save_group(store, *args, **kwargs):
203
208
raise ValueError ('at least one array must be provided' )
204
209
# handle polymorphic store arg
205
210
may_need_closing = isinstance (store , str )
206
- store = normalize_store_arg (store , clobber = True )
211
+ _store : Store = normalize_store_arg (store , clobber = True )
207
212
try :
208
- grp = _create_group (store , overwrite = True )
213
+ grp = _create_group (_store , overwrite = True )
209
214
for i , arr in enumerate (args ):
210
215
k = 'arr_{}' .format (i )
211
216
grp .create_dataset (k , data = arr , overwrite = True )
212
217
for k , arr in kwargs .items ():
213
218
grp .create_dataset (k , data = arr , overwrite = True )
214
219
finally :
215
- if may_need_closing and hasattr ( store , 'close' ) :
220
+ if may_need_closing :
216
221
# needed to ensure zip file records are written
217
- store .close ()
222
+ _store .close ()
218
223
219
224
220
- def save (store , * args , ** kwargs ):
225
+ def save (store : StoreLike , * args , ** kwargs ):
221
226
"""Convenience function to save an array or group of arrays to the local file system.
222
227
223
228
Parameters
@@ -327,7 +332,7 @@ def __repr__(self):
327
332
return r
328
333
329
334
330
- def load (store ):
335
+ def load (store : StoreLike ):
331
336
"""Load data from an array or group into memory.
332
337
333
338
Parameters
@@ -353,11 +358,11 @@ def load(store):
353
358
354
359
"""
355
360
# handle polymorphic store arg
356
- store = normalize_store_arg (store )
357
- if contains_array (store , path = None ):
358
- return Array (store = store , path = None )[...]
359
- elif contains_group (store , path = None ):
360
- grp = Group (store = store , path = None )
361
+ _store = normalize_store_arg (store )
362
+ if contains_array (_store , path = None ):
363
+ return Array (store = _store , path = None )[...]
364
+ elif contains_group (_store , path = None ):
365
+ grp = Group (store = _store , path = None )
361
366
return LazyLoader (grp )
362
367
363
368
@@ -1073,7 +1078,7 @@ def copy_all(source, dest, shallow=False, without_attrs=False, log=None,
1073
1078
return n_copied , n_skipped , n_bytes_copied
1074
1079
1075
1080
1076
- def consolidate_metadata (store , metadata_key = ' .zmetadata' ):
1081
+ def consolidate_metadata (store : Store , metadata_key = " .zmetadata" ):
1077
1082
"""
1078
1083
Consolidate all metadata for groups and arrays within the given store
1079
1084
into a single resource and put it under the given key.
@@ -1124,7 +1129,7 @@ def is_zarr_key(key):
1124
1129
return open_consolidated (store , metadata_key = metadata_key )
1125
1130
1126
1131
1127
- def open_consolidated (store , metadata_key = ' .zmetadata' , mode = 'r+' , ** kwargs ):
1132
+ def open_consolidated (store : Store , metadata_key = " .zmetadata" , mode = "r+" , ** kwargs ):
1128
1133
"""Open group using metadata previously consolidated into a single key.
1129
1134
1130
1135
This is an optimised method for opening a Zarr group, where instead of
0 commit comments