8
8
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
9
9
# module imports
10
10
"""Utilities to load and save image objects"""
11
+ from __future__ import annotations
12
+
11
13
import os
14
+ import typing as ty
12
15
13
16
import numpy as np
14
17
22
25
_compressed_suffixes = ('.gz' , '.bz2' , '.zst' )
23
26
24
27
25
- def _signature_matches_extension (filename ):
28
+ if ty .TYPE_CHECKING : # pragma: no cover
29
+ from .filebasedimages import FileBasedImage
30
+ from .filename_parser import FileSpec
31
+
32
+ P = ty .ParamSpec ('P' )
33
+
34
+ class Signature (ty .TypedDict ):
35
+ signature : bytes
36
+ format_name : str
37
+
38
+
39
+ def _signature_matches_extension (filename : FileSpec ) -> tuple [bool , str ]:
26
40
"""Check if signature aka magic number matches filename extension.
27
41
28
42
Parameters
@@ -42,7 +56,7 @@ def _signature_matches_extension(filename):
42
56
the empty string otherwise.
43
57
44
58
"""
45
- signatures = {
59
+ signatures : dict [ str , Signature ] = {
46
60
'.gz' : {'signature' : b'\x1f \x8b ' , 'format_name' : 'gzip' },
47
61
'.bz2' : {'signature' : b'BZh' , 'format_name' : 'bzip2' },
48
62
'.zst' : {'signature' : b'\x28 \xb5 \x2f \xfd ' , 'format_name' : 'ztsd' },
@@ -64,7 +78,7 @@ def _signature_matches_extension(filename):
64
78
return False , f'File { filename } is not a { format_name } file'
65
79
66
80
67
- def load (filename , ** kwargs ):
81
+ def load (filename : FileSpec , ** kwargs ) -> FileBasedImage :
68
82
r"""Load file given filename, guessing at file type
69
83
70
84
Parameters
@@ -126,7 +140,7 @@ def guessed_image_type(filename):
126
140
raise ImageFileError (f'Cannot work out file type of "{ filename } "' )
127
141
128
142
129
- def save (img , filename , ** kwargs ):
143
+ def save (img : FileBasedImage , filename : FileSpec , ** kwargs ) -> None :
130
144
r"""Save an image to file adapting format to `filename`
131
145
132
146
Parameters
@@ -161,19 +175,17 @@ def save(img, filename, **kwargs):
161
175
from .nifti1 import Nifti1Image , Nifti1Pair
162
176
from .nifti2 import Nifti2Image , Nifti2Pair
163
177
164
- klass = None
165
- converted = None
166
-
178
+ converted : FileBasedImage
167
179
if type (img ) == Nifti1Image and lext in ('.img' , '.hdr' ):
168
- klass = Nifti1Pair
180
+ converted = Nifti1Pair . from_image ( img )
169
181
elif type (img ) == Nifti2Image and lext in ('.img' , '.hdr' ):
170
- klass = Nifti2Pair
182
+ converted = Nifti2Pair . from_image ( img )
171
183
elif type (img ) == Nifti1Pair and lext == '.nii' :
172
- klass = Nifti1Image
184
+ converted = Nifti1Image . from_image ( img )
173
185
elif type (img ) == Nifti2Pair and lext == '.nii' :
174
- klass = Nifti2Image
186
+ converted = Nifti2Image . from_image ( img )
175
187
else : # arbitrary conversion
176
- valid_klasses = [klass for klass in all_image_classes if ext in klass .valid_exts ]
188
+ valid_klasses = [klass for klass in all_image_classes if lext in klass .valid_exts ]
177
189
if not valid_klasses : # if list is empty
178
190
raise ImageFileError (f'Cannot work out file type of "{ filename } "' )
179
191
@@ -186,13 +198,9 @@ def save(img, filename, **kwargs):
186
198
break
187
199
except Exception as e :
188
200
err = e
189
- # ... and if none of them work, raise an error.
190
- if converted is None :
201
+ else :
191
202
raise err
192
203
193
- # Here, we either have a klass or a converted image.
194
- if converted is None :
195
- converted = klass .from_image (img )
196
204
converted .to_filename (filename , ** kwargs )
197
205
198
206
0 commit comments