Enhance HDF5 attribute handling

Author: kojix2

README.md

@@ -11,12 +11,11 @@ This gem currently provides practical high-level wrappers for:
 - opening and creating files
 - creating groups
 - creating, writing, and reading one-dimensional numeric datasets
-- reading attributes
+- reading and writing numeric attributes
 
 Unsupported at this stage:
 
 - string dataset read/write
-- attribute write
 - multidimensional array write
 
 ## Supported HDF5 Versions

lib/hdf5/attribute.rb

@@ -62,5 +62,70 @@ def [](attr_name)
     ensure
       attr.close if attr
     end
+
+    def []=(attr_name, value)
+      write(attr_name, value)
+    end
+
+    def write(attr_name, value)
+      values = normalize_data(value)
+      datatype_id = datatype_id_for(values)
+
+      dims = ::FFI::MemoryPointer.new(:ulong_long, 1)
+      dims.write_array_of_ulong_long([values.length])
+      dataspace_id = HDF5::FFI.H5Screate_simple(1, dims, nil)
+      raise HDF5::Error, 'Failed to create attribute dataspace' if dataspace_id < 0
+
+      attr_id = HDF5::FFI.H5Acreate2(
+        @dataset_id,
+        attr_name,
+        datatype_id,
+        dataspace_id,
+        HDF5::DEFAULT_PROPERTY_LIST,
+        HDF5::DEFAULT_PROPERTY_LIST
+      )
+      raise HDF5::Error, "Failed to create attribute: #{attr_name}" if attr_id < 0
+
+      buffer = buffer_for(values)
+      status = HDF5::FFI.H5Awrite(attr_id, datatype_id, buffer)
+      raise HDF5::Error, "Failed to write attribute: #{attr_name}" if status < 0
+
+      value
+    ensure
+      HDF5::FFI.H5Aclose(attr_id) if attr_id && attr_id >= 0
+      HDF5::FFI.H5Sclose(dataspace_id) if dataspace_id && dataspace_id >= 0
+    end
+
+    private
+
+    def normalize_data(value)
+      values = value.is_a?(Array) ? value : [value]
+      raise HDF5::Error, 'Attribute data must not be empty' if values.empty?
+      raise HDF5::Error, 'Nested arrays are not supported for attributes' if values.any? { |item| item.is_a?(Array) }
+
+      values
+    end
+
+    def datatype_id_for(values)
+      if values.all? { |item| item.is_a?(Integer) }
+        HDF5::FFI.H5T_NATIVE_INT
+      elsif values.all? { |item| item.is_a?(Numeric) }
+        HDF5::FFI.H5T_NATIVE_DOUBLE
+      else
+        raise HDF5::Error, 'Only numeric attribute data is supported'
+      end
+    end
+
+    def buffer_for(values)
+      if values.all? { |item| item.is_a?(Integer) }
+        buffer = ::FFI::MemoryPointer.new(:int, values.length)
+        buffer.write_array_of_int(values)
+      else
+        buffer = ::FFI::MemoryPointer.new(:double, values.length)
+        buffer.write_array_of_double(values.map(&:to_f))
+      end
+
+      buffer
+    end
   end
 end

lib/hdf5/dataset.rb

@@ -66,6 +66,8 @@ def buffer_for(data)
         buffer
       end
 
+      private :normalize_data, :datatype_id_for, :buffer_for
+
       private
 
       def from_id(dataset_id, name)
@@ -84,9 +86,9 @@ def attrs
     end
 
     def write(data)
-      values = self.class.normalize_data(data)
-      buffer = self.class.buffer_for(values)
-      mem_type_id = self.class.datatype_id_for(values)
+      values = self.class.send(:normalize_data, data)
+      buffer = self.class.send(:buffer_for, values)
+      mem_type_id = self.class.send(:datatype_id_for, values)
       status = HDF5::FFI.H5Dwrite(@dataset_id, mem_type_id, HDF5::DEFAULT_PROPERTY_LIST, HDF5::DEFAULT_PROPERTY_LIST,
                                   HDF5::DEFAULT_PROPERTY_LIST, buffer)
       raise HDF5::Error, 'Failed to write dataset' if status < 0
@@ -160,8 +162,8 @@ def read_float_data(total_elements)
       buffer.read_array_of_double(total_elements)
     end
 
-    def read_string_data(total_elements)
-      raise NotImplementedError
+    def read_string_data(_total_elements)
+      raise HDF5::Error, 'String dataset reading is not supported yet'
     end
 
     private

lib/hdf5/group.rb

@@ -54,10 +54,10 @@ def create_dataset(name, data, &block)
       Dataset.create(@group_id, name, data, &block)
     end
 
-    def list_datasets
-      datasets = []
+    def list_entries
+      entries = []
       callback = ::FFI::Function.new(:int, %i[int64_t string pointer pointer]) do |_, name, _, _|
-        datasets << name
+        entries << name
         0 # continue
       end
 
@@ -66,9 +66,13 @@ def list_datasets
        else
          HDF5::FFI.H5Literate2(@group_id, :H5_INDEX_NAME, :H5_ITER_NATIVE, nil, callback, nil)
        end).negative? &&
-        raise(HDF5::Error, 'Failed to list datasets')
+        raise(HDF5::Error, 'Failed to list entries')
+
+      entries
+    end
 
-      datasets
+    def list_datasets
+      list_entries.select { |name| dataset?(name) }
     end
 
     def [](name)

test/hdf5_test.rb

@@ -96,27 +96,7 @@ class HDF5Test < Test::Unit::TestCase
 
       HDF5::File.create(path) do |file|
         dataset = file.create_dataset('values', [1, 2, 3])
-        dataset_id = dataset.instance_variable_get(:@dataset_id)
-
-        dims = FFI::MemoryPointer.new(:ulong_long, 1)
-        dims.write_array_of_ulong_long([1])
-        attr_space_id = HDF5::FFI.H5Screate_simple(1, dims, nil)
-        attr_id = HDF5::FFI.H5Acreate2(
-          dataset_id,
-          'scale',
-          HDF5::FFI.H5T_NATIVE_INT,
-          attr_space_id,
-          HDF5::DEFAULT_PROPERTY_LIST,
-          HDF5::DEFAULT_PROPERTY_LIST
-        )
-        value = FFI::MemoryPointer.new(:int, 1)
-        value.write_int(42)
-        status = HDF5::FFI.H5Awrite(attr_id, HDF5::FFI.H5T_NATIVE_INT, value)
-        assert_equal(0, status)
-      ensure
-        HDF5::FFI.H5Aclose(attr_id) if attr_id && attr_id >= 0
-        HDF5::FFI.H5Sclose(attr_space_id) if attr_space_id && attr_space_id >= 0
-        dataset.close if dataset
+        dataset.attrs['scale'] = 42
       end
 
       HDF5::File.open(path) do |file|
@@ -125,6 +105,25 @@ class HDF5Test < Test::Unit::TestCase
     end
   end
 
+  test 'group list_datasets returns only datasets' do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, 'group-list.h5')
+
+      HDF5::File.create(path) do |file|
+        file.create_group('parent') do |parent|
+          parent.create_group('child')
+          parent.create_dataset('numbers', [1, 2, 3])
+        end
+      end
+
+      HDF5::File.open(path) do |file|
+        parent = file['parent']
+        assert_equal(%w[child numbers].sort, parent.list_entries.sort)
+        assert_equal(%w[numbers], parent.list_datasets)
+      end
+    end
+  end
+
   test 'dataset attrs raises HDF5 error for missing attribute' do
     Dir.mktmpdir do |dir|
       path = File.join(dir, 'missing-attribute.h5')