ENH: Enforce GIFTI compatibility at write

Author: effigies

nibabel/gifti/gifti.py

@@ -16,7 +16,8 @@
 import base64
 import sys
 import warnings
-from typing import Type
+from copy import copy
+from typing import Type, cast
 
 import numpy as np
 
@@ -27,6 +28,12 @@
 from ..nifti1 import data_type_codes, intent_codes, xform_codes
 from .util import KIND2FMT, array_index_order_codes, gifti_encoding_codes, gifti_endian_codes
 
+GIFTI_DTYPES = (
+    data_type_codes['NIFTI_TYPE_UINT8'],
+    data_type_codes['NIFTI_TYPE_INT32'],
+    data_type_codes['NIFTI_TYPE_FLOAT32'],
+)
+
 
 class _GiftiMDList(list):
     """List view of GiftiMetaData object that will translate most operations"""
@@ -462,11 +469,7 @@ def __init__(
         if datatype is None:
             if self.data is None:
                 datatype = 'none'
-            elif self.data.dtype in (
-                np.dtype('uint8'),
-                np.dtype('int32'),
-                np.dtype('float32'),
-            ):
+            elif data_type_codes[self.data.dtype] in GIFTI_DTYPES:
                 datatype = self.data.dtype
             else:
                 raise ValueError(
@@ -848,20 +851,45 @@ def _to_xml_element(self):
             GIFTI.append(dar._to_xml_element())
         return GIFTI
 
-    def to_xml(self, enc='utf-8') -> bytes:
+    def to_xml(self, enc='utf-8', *, mode='strict') -> bytes:
         """Return XML corresponding to image content"""
+        if mode == 'strict':
+            if any(arr.datatype not in GIFTI_DTYPES for arr in self.darrays):
+                raise ValueError(
+                    'GiftiImage contains data arrays with invalid data types; '
+                    'use mode="compat" to automatically cast to conforming types'
+                )
+        elif mode == 'compat':
+            darrays = []
+            for arr in self.darrays:
+                if arr.datatype not in GIFTI_DTYPES:
+                    arr = copy(arr)
+                    # TODO: Better typing for recoders
+                    dtype = cast(np.dtype, data_type_codes.dtype[arr.datatype])
+                    if np.issubdtype(dtype, np.floating):
+                        arr.datatype = data_type_codes['float32']
+                    elif np.issubdtype(dtype, np.integer):
+                        arr.datatype = data_type_codes['int32']
+                    else:
+                        raise ValueError(f'Cannot convert {dtype} to float32/int32')
+                darrays.append(arr)
+            gii = copy(self)
+            gii.darrays = darrays
+            return gii.to_xml(enc=enc, mode='strict')
+        elif mode != 'force':
+            raise TypeError(f'Unknown mode {mode}')
         header = b"""<?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE GIFTI SYSTEM "http://www.nitrc.org/frs/download.php/115/gifti.dtd">
 """
         return header + super().to_xml(enc)
 
     # Avoid the indirection of going through to_file_map
-    def to_bytes(self, enc='utf-8'):
-        return self.to_xml(enc=enc)
+    def to_bytes(self, enc='utf-8', *, mode='strict'):
+        return self.to_xml(enc=enc, mode=mode)
 
     to_bytes.__doc__ = SerializableImage.to_bytes.__doc__
 
-    def to_file_map(self, file_map=None, enc='utf-8'):
+    def to_file_map(self, file_map=None, enc='utf-8', *, mode='strict'):
         """Save the current image to the specified file_map
 
         Parameters
@@ -877,7 +905,7 @@ def to_file_map(self, file_map=None, enc='utf-8'):
         if file_map is None:
             file_map = self.file_map
         with file_map['image'].get_prepare_fileobj('wb') as f:
-            f.write(self.to_xml(enc=enc))
+            f.write(self.to_xml(enc=enc, mode=mode))
 
     @classmethod
     def from_file_map(klass, file_map, buffer_size=35000000, mmap=True):

nibabel/gifti/tests/test_gifti.py

@@ -505,7 +505,7 @@ def test_darray_dtype_coercion_failures():
             datatype=darray_dtype,
         )
         gii = GiftiImage(darrays=[da])
-        gii_copy = GiftiImage.from_bytes(gii.to_bytes())
+        gii_copy = GiftiImage.from_bytes(gii.to_bytes(mode='force'))
         da_copy = gii_copy.darrays[0]
         assert np.dtype(da_copy.data.dtype) == np.dtype(darray_dtype)
         assert_array_equal(da_copy.data, da.data)