44from functools import partial
55from pathlib import Path
66from types import MappingProxyType
7- from typing import TYPE_CHECKING , TypeVar
7+ from typing import TYPE_CHECKING , TypeVar , cast
88from warnings import warn
99
1010import h5py
3636)
3737
3838if TYPE_CHECKING :
39- from collections .abc import Callable , Collection , Mapping , Sequence
39+ from collections .abc import Callable , Collection , Container , Mapping , Sequence
4040 from os import PathLike
4141 from typing import Any , Literal
4242
4343 from .._core .file_backing import AnnDataFileManager
44+ from .._core .raw import Raw
4445
4546T = TypeVar ("T" )
4647
@@ -82,29 +83,18 @@ def write_h5ad(
8283 # TODO: Use spec writing system for this
8384 # Currently can't use write_dispatched here because this function is also called to do an
8485 # inplace update of a backed object, which would delete "/"
85- f = f ["/" ]
86+ f = cast ( "h5py.Group" , f ["/" ])
8687 f .attrs .setdefault ("encoding-type" , "anndata" )
8788 f .attrs .setdefault ("encoding-version" , "0.1.0" )
8889
89- if "X" in as_dense and isinstance (
90- adata .X , CSMatrix | BaseCompressedSparseDataset
91- ):
92- write_sparse_as_dense (f , "X" , adata .X , dataset_kwargs = dataset_kwargs )
93- elif not (adata .isbacked and Path (adata .filename ) == Path (filepath )):
94- # If adata.isbacked, X should already be up to date
95- write_elem (f , "X" , adata .X , dataset_kwargs = dataset_kwargs )
96- if "raw/X" in as_dense and isinstance (
97- adata .raw .X , CSMatrix | BaseCompressedSparseDataset
98- ):
99- write_sparse_as_dense (
100- f , "raw/X" , adata .raw .X , dataset_kwargs = dataset_kwargs
101- )
102- write_elem (f , "raw/var" , adata .raw .var , dataset_kwargs = dataset_kwargs )
103- write_elem (
104- f , "raw/varm" , dict (adata .raw .varm ), dataset_kwargs = dataset_kwargs
105- )
106- elif adata .raw is not None :
107- write_elem (f , "raw" , adata .raw , dataset_kwargs = dataset_kwargs )
90+ _write_x (
91+ f ,
92+ adata , # accessing adata.X reopens adata.file if it’s backed
93+ is_backed = adata .isbacked and adata .filename == filepath ,
94+ as_dense = as_dense ,
95+ dataset_kwargs = dataset_kwargs ,
96+ )
97+ _write_raw (f , adata .raw , as_dense = as_dense , dataset_kwargs = dataset_kwargs )
10898 write_elem (f , "obs" , adata .obs , dataset_kwargs = dataset_kwargs )
10999 write_elem (f , "var" , adata .var , dataset_kwargs = dataset_kwargs )
110100 write_elem (f , "obsm" , dict (adata .obsm ), dataset_kwargs = dataset_kwargs )
@@ -115,6 +105,41 @@ def write_h5ad(
115105 write_elem (f , "uns" , dict (adata .uns ), dataset_kwargs = dataset_kwargs )
116106
117107
108+ def _write_x (
109+ f : h5py .Group ,
110+ adata : AnnData ,
111+ * ,
112+ is_backed : bool ,
113+ as_dense : Container [str ],
114+ dataset_kwargs : Mapping [str , Any ],
115+ ) -> None :
116+ if "X" in as_dense and isinstance (adata .X , CSMatrix | BaseCompressedSparseDataset ):
117+ write_sparse_as_dense (f , "X" , adata .X , dataset_kwargs = dataset_kwargs )
118+ elif is_backed :
119+ pass # If adata.isbacked, X should already be up to date
120+ elif adata .X is None :
121+ f .pop ("X" , None )
122+ else :
123+ write_elem (f , "X" , adata .X , dataset_kwargs = dataset_kwargs )
124+
125+
126+ def _write_raw (
127+ f : h5py .Group ,
128+ raw : Raw ,
129+ * ,
130+ as_dense : Container [str ],
131+ dataset_kwargs : Mapping [str , Any ],
132+ ) -> None :
133+ if "raw/X" in as_dense and isinstance (
134+ raw .X , CSMatrix | BaseCompressedSparseDataset
135+ ):
136+ write_sparse_as_dense (f , "raw/X" , raw .X , dataset_kwargs = dataset_kwargs )
137+ write_elem (f , "raw/var" , raw .var , dataset_kwargs = dataset_kwargs )
138+ write_elem (f , "raw/varm" , dict (raw .varm ), dataset_kwargs = dataset_kwargs )
139+ elif raw is not None :
140+ write_elem (f , "raw" , raw , dataset_kwargs = dataset_kwargs )
141+
142+
118143@report_write_key_on_error
119144@write_spec (IOSpec ("array" , "0.2.0" ))
120145def write_sparse_as_dense (
0 commit comments