1717from typing import (
1818 TYPE_CHECKING ,
1919 Any ,
20- Final ,
2120 Literal ,
2221 TypeVar ,
2322 Union ,
9897DATAARRAY_NAME = "__xarray_dataarray_name__"
9998DATAARRAY_VARIABLE = "__xarray_dataarray_variable__"
10099
101- ENGINES = {
102- "netcdf4" : backends .NetCDF4DataStore .open ,
103- "scipy" : backends .ScipyDataStore ,
104- "pydap" : backends .PydapDataStore .open ,
105- "h5netcdf" : backends .H5NetCDFStore .open ,
106- "zarr" : backends .ZarrStore .open_group ,
107- }
108-
109-
110- def _get_default_engine_remote_uri () -> Literal ["netcdf4" , "pydap" ]:
111- engine : Literal ["netcdf4" , "pydap" ]
112- try :
113- import netCDF4 # noqa: F401
114-
115- engine = "netcdf4"
116- except ImportError : # pragma: no cover
117- try :
118- import pydap # noqa: F401
119100
120- engine = "pydap"
121- except ImportError as err :
122- raise ValueError (
123- "netCDF4 or pydap is required for accessing remote datasets via OPeNDAP"
124- ) from err
125- return engine
126-
127-
128- def _get_default_engine_gz () -> Literal ["scipy" ]:
129- try :
130- import scipy # noqa: F401
101+ def get_default_netcdf_write_engine (
102+ format : T_NetcdfTypes | None ,
103+ to_fileobject_or_memoryview : bool ,
104+ ) -> Literal ["netcdf4" , "h5netcdf" , "scipy" ]:
105+ """Return the default netCDF library to use for writing a netCDF file."""
106+ module_names = {
107+ "netcdf4" : "netCDF4" ,
108+ "scipy" : "scipy" ,
109+ "h5netcdf" : "h5netcdf" ,
110+ }
131111
132- engine : Final = "scipy"
133- except ImportError as err : # pragma: no cover
134- raise ValueError ("scipy is required for accessing .gz files" ) from err
135- return engine
112+ candidates = list (plugins .NETCDF_BACKENDS_ORDER )
136113
114+ if format is not None :
115+ if format .upper ().startswith ("NETCDF3" ):
116+ candidates .remove ("h5netcdf" )
117+ elif format .upper ().startswith ("NETCDF4" ):
118+ candidates .remove ("scipy" )
119+ else :
120+ raise ValueError (f"unexpected { format = } )
137121
138- def _get_default_engine_netcdf () -> Literal ["netcdf4" , "h5netcdf" , "scipy" ]:
139- candidates : list [tuple [str , str ]] = [
140- ("netcdf4" , "netCDF4" ),
141- ("h5netcdf" , "h5netcdf" ),
142- ("scipy" , "scipy.io.netcdf" ),
143- ]
122+ if to_fileobject_or_memoryview :
123+ candidates .remove ("netcdf4" )
144124
145- for engine , module_name in candidates :
125+ for engine in candidates :
126+ module_name = module_names [engine ]
146127 if importlib .util .find_spec (module_name ) is not None :
147128 return cast (Literal ["netcdf4" , "h5netcdf" , "scipy" ], engine )
148129
130+ format_str = f" with { format = } if format is not None else ""
131+ libraries = ", " .join (module_names [c ] for c in candidates )
149132 raise ValueError (
150- "cannot read or write NetCDF files because none of "
151- "'netCDF4-python', 'h5netcdf', or 'scipy' are installed"
133+ f "cannot write NetCDF files{ format_str } the suitable "
134+ f"backend libraries ( { libraries } ) are installed"
152135 )
153136
154137
155- def _get_default_engine (path : str , allow_remote : bool = False ) -> T_NetcdfEngine :
156- if allow_remote and is_remote_uri (path ):
157- return _get_default_engine_remote_uri () # type: ignore[return-value]
158- elif path .endswith (".gz" ):
159- return _get_default_engine_gz ()
160- else :
161- return _get_default_engine_netcdf ()
162-
163-
164138def _validate_dataset_names (dataset : Dataset ) -> None :
165139 """DataArray.name and Dataset keys must be a string or None"""
166140
@@ -1958,7 +1932,7 @@ def to_netcdf(
19581932 multifile : Literal [False ] = False ,
19591933 invalid_netcdf : bool = False ,
19601934 auto_complex : bool | None = None ,
1961- ) -> bytes | memoryview : ...
1935+ ) -> memoryview : ...
19621936
19631937
19641938# compute=False returns dask.Delayed
@@ -2051,7 +2025,7 @@ def to_netcdf(
20512025 multifile : bool = False ,
20522026 invalid_netcdf : bool = False ,
20532027 auto_complex : bool | None = None ,
2054- ) -> tuple [ArrayWriter , AbstractDataStore ] | bytes | memoryview | Delayed | None : ...
2028+ ) -> tuple [ArrayWriter , AbstractDataStore ] | memoryview | Delayed | None : ...
20552029
20562030
20572031def to_netcdf (
@@ -2067,41 +2041,22 @@ def to_netcdf(
20672041 multifile : bool = False ,
20682042 invalid_netcdf : bool = False ,
20692043 auto_complex : bool | None = None ,
2070- ) -> tuple [ArrayWriter , AbstractDataStore ] | bytes | memoryview | Delayed | None :
2044+ ) -> tuple [ArrayWriter , AbstractDataStore ] | memoryview | Delayed | None :
20712045 """This function creates an appropriate datastore for writing a dataset to
20722046 disk as a netCDF file
20732047
20742048 See `Dataset.to_netcdf` for full API docs.
20752049
20762050 The ``multifile`` argument is only for the private use of save_mfdataset.
20772051 """
2078- if isinstance (path_or_file , os .PathLike ):
2079- path_or_file = os .fspath (path_or_file )
2080-
20812052 if encoding is None :
20822053 encoding = {}
20832054
2084- if isinstance (path_or_file , str ):
2085- if engine is None :
2086- engine = _get_default_engine (path_or_file )
2087- path_or_file = _normalize_path (path_or_file )
2088- else :
2089- # writing to bytes/memoryview or a file-like object
2090- if engine is None :
2091- # TODO: only use 'scipy' if format is None or a netCDF3 format
2092- engine = "scipy"
2093- elif engine not in "scipy" , "h5netcdf" ):
2094- raise ValueError (
2095- "invalid engine for creating bytes/memoryview or writing to a "
2096- f"file-like object with to_netcdf: { engine !r}
2097- "engine=None, engine='scipy' and engine='h5netcdf' is "
2098- "supported."
2099- )
2100- if not compute :
2101- raise NotImplementedError (
2102- "to_netcdf() with compute=False is not yet implemented when "
2103- "returning bytes"
2104- )
2055+ path_or_file = _normalize_path (path_or_file )
2056+
2057+ if engine is None :
2058+ to_fileobject_or_memoryview = not isinstance (path_or_file , str )
2059+ engine = get_default_netcdf_write_engine (format , to_fileobject_or_memoryview )
21052060
21062061 # validate Dataset keys, DataArray names, and attr keys/values
21072062 _validate_dataset_names (dataset )
@@ -2121,6 +2076,11 @@ def to_netcdf(
21212076 )
21222077
21232078 if path_or_file is None :
2079+ if not compute :
2080+ raise NotImplementedError (
2081+ "to_netcdf() with compute=False is not yet implemented when "
2082+ "returning a memoryview"
2083+ )
21242084 target = BytesIOProxy ()
21252085 else :
21262086 target = path_or_file # type: ignore[assignment]
@@ -2164,7 +2124,7 @@ def to_netcdf(
21642124
21652125 if path_or_file is None :
21662126 assert isinstance (target , BytesIOProxy ) # created in this function
2167- return target .getvalue_or_getbuffer ()
2127+ return target .getbuffer ()
21682128
21692129 if not compute :
21702130 return delayed_close_after_writes (writes , store )
0 commit comments