xopr.matlab_attribute_utils

  1"""
  2Utilities and fixes for loading MATLAB format files into Xarray.
  3
  4This module provides various utilities for correctly reading MATLAB files and
  5converting them to formats compatible with xarray datasets. It handles both
  6modern HDF5-format MATLAB files (.mat v7.3+) and legacy MATLAB file formats.
  7
  8The module addresses several common issues when loading MATLAB data:
  9
 10- Dereferencing HDF5 object references in MATLAB files
 11- Decoding MATLAB char arrays (both uint16 Unicode and uint8 ASCII)
 12- Converting MATLAB cell arrays to Python lists
 13- Handling empty MATLAB arrays
 14- Stripping sensitive data (API keys) from attributes
 15- Converting object ndarrays to JSON-serializable lists
 16
 17HDF5-Format MATLAB Files (v7.3+)
 18--------------------------------
 19- dereference_h5value: Recursively dereference HDF5 object references
 20- decode_hdf5_matlab_variable: Decode MATLAB variables from HDF5 storage
 21
 22Legacy MATLAB Files (v4-v7.2)
 23------------------------------
 24- extract_legacy_mat_attributes: Extract attributes from legacy .mat files
 25- strip_api_key: Remove API keys from attribute dictionaries
 26- convert_object_ndarrays_to_lists: Convert object arrays to lists
 27
 28Notes
 29-----
 30@private This module is not intended for external use.
 31
 32"""
 33
 34from collections.abc import Iterable
 35
 36import h5py
 37import numpy as np
 38import scipy.io
 39
 40#
 41# HDF5-format MATLAB files
 42#
 43
 44def dereference_h5value(value, h5file, make_array=True):
 45    if isinstance(value, h5py.Reference):
 46        return dereference_h5value(h5file[value], h5file=h5file)
 47    elif isinstance(value, h5py.Group):
 48        # Pass back to decode_hdf5_matlab_variable to handle groups
 49        return decode_hdf5_matlab_variable(value, h5file=h5file)
 50    elif isinstance(value, Iterable):
 51        v = [dereference_h5value(v, h5file=h5file) for v in value]
 52        if make_array:
 53            try:
 54                return np.squeeze(np.array(v))
 55            except:
 56                return v
 57        else:
 58            return v
 59    elif isinstance(value, np.number):
 60        return value.item()
 61    else:
 62        return value
 63
 64def decode_hdf5_matlab_variable(h5var, skip_variables=False, debug_path="", skip_errors=True, h5file=None):
 65    """
 66    Decode a MATLAB variable stored in an HDF5 file.
 67    This function assumes the variable is stored as a byte string.
 68    """
 69    if h5file is None:
 70        h5file = h5var.file
 71    matlab_class = h5var.attrs.get('MATLAB_class', None)
 72
 73    # Handle MATLAB_class as either bytes or string
 74    if matlab_class and (matlab_class == b'cell' or matlab_class == 'cell'):
 75        return dereference_h5value(h5var[:], h5file=h5file, make_array=False)
 76    elif matlab_class and (matlab_class == b'char' or matlab_class == 'char'):
 77        # Check if this is an empty MATLAB char array
 78        if h5var.attrs.get('MATLAB_empty', 0):
 79            return ''
 80
 81        # MATLAB stores char arrays as uint16 (Unicode code points)
 82        # or sometimes uint8 (ASCII). Handle both cases properly.
 83        data = h5var[:]
 84
 85        if data.dtype == np.dtype('uint16'):
 86            # Each uint16 value is a Unicode code point (UCS-2/UTF-16)
 87            # Convert to string by treating each value as a character code
 88            chars = [chr(c) for c in data.flatten() if c != 0]
 89            return ''.join(chars).rstrip()
 90        elif data.dtype == np.dtype('uint8'):
 91            # uint8 data can be decoded directly as UTF-8
 92            return data.tobytes().decode('utf-8').rstrip('\x00')
 93        else:
 94            # Fallback for unexpected dtypes (including uint64 for empty arrays)
 95            # First check if it's all zeros (empty string)
 96            if np.all(data == 0):
 97                return ''
 98            # Try the old method that may work for some cases
 99            try:
100                return data.astype(dtype=np.uint8).tobytes().decode('utf-8').rstrip('\x00')
101            except UnicodeDecodeError:
102                # If that fails, try to convert assuming Unicode code points
103                chars = [chr(min(c, 0x10FFFF)) for c in data.flatten() if c != 0]
104                return ''.join(chars).rstrip()
105    elif isinstance(h5var, (h5py.Group, h5py.File)):
106        attrs = {}
107        for k in h5var:
108            if k.startswith('#'):
109                continue
110            if 'api_key' in k:
111                attrs[k] = "API_KEY_REMOVED"
112                continue
113            if isinstance(h5var[k], h5py.Dataset):
114                if not skip_variables:
115                    try:
116                        attrs[k] = decode_hdf5_matlab_variable(h5var[k], debug_path=debug_path + "/" + k, skip_errors=skip_errors, h5file=h5file)
117                    except Exception as e:
118                        print(f"Failed to decode variable {k} at {debug_path}: {e}")
119                        if not skip_errors:
120                            raise e
121            else:
122                attrs[k] = decode_hdf5_matlab_variable(h5var[k], debug_path=debug_path + "/" + k, skip_errors=skip_errors, h5file=h5file)
123        return attrs
124    elif isinstance(h5var, h5py.Dataset):
125        if h5var.dtype == 'O':
126            return dereference_h5value(h5var[:], h5file=h5file)
127        else:
128            return np.squeeze(h5var[:])
129    else:
130        return h5var[:]
131
132#
133# Legacy MATLAB files (non-HDF5)
134#
135
136def extract_legacy_mat_attributes(file, skip_keys=[], skip_errors=True):
137    m = scipy.io.loadmat(file, mat_dtype=False, simplify_cells=True, squeeze_me=True)
138
139    attrs = {key: value for key, value in m.items()
140             if not key.startswith('__') and key not in skip_keys}
141
142    attrs = strip_api_key(attrs)
143    attrs = convert_object_ndarrays_to_lists(attrs)
144    return attrs
145
146def strip_api_key(attrs):
147    attrs_clean = {}
148    for key, value in attrs.items():
149        if 'api_key' in key:
150            attrs_clean[key] = "API_KEY_REMOVED"
151        elif isinstance(value, dict):
152            attrs_clean[key] = strip_api_key(value)
153        else:
154            attrs_clean[key] = value
155    return attrs_clean
156
157def convert_object_ndarrays_to_lists(attrs):
158    """
159    Convert any object ndarray attributes to lists.
160    """
161    for key, value in attrs.items():
162        if isinstance(value, np.ndarray) and value.dtype == 'object':
163            attrs[key] = value.tolist()
164        elif isinstance(value, dict):
165            convert_object_ndarrays_to_lists(value)
166        else:
167            attrs[key] = value
168    return attrs