Fix npz loading

minhqdao · minhqdao · commit 25fdd87ba169 · 2024-02-23T03:13:42.000+03:00
diff --git a/src/stdlib_array.fypp b/src/stdlib_array.fypp
@@ -14,8 +14,18 @@ module stdlib_array
 
     public :: trueloc, falseloc
 
-    type, public :: t_array_bundle
-        class(t_array), allocatable :: files(:)
+    !> Helper class to allocate t_array as an abstract type.
+    type, public :: t_array_wrapper
+        class(t_array), allocatable :: array
+
+    contains
+
+        #:for k1, t1 in KINDS_TYPES
+            #:for rank in RANKS
+                generic :: allocate_array => allocate_array_${t1[0]}$${k1}$_${rank}$
+                procedure :: allocate_array_${t1[0]}$${k1}$_${rank}$
+            #:endfor
+        #:endfor
     end type
 
     type, abstract, public :: t_array
@@ -32,6 +42,30 @@ module stdlib_array
 
 contains
 
+    #:for k1, t1 in KINDS_TYPES
+        #:for rank in RANKS
+            !> Allocate an instance of the array within the wrapper.
+            module subroutine allocate_array_${t1[0]}$${k1}$_${rank}$ (wrapper, array, stat, msg)
+                class(t_array_wrapper), intent(out) :: wrapper
+                ${t1}$, intent(in) :: array${ranksuffix(rank)}$
+                integer, intent(out) :: stat
+                character(len=:), allocatable, intent(out) :: msg
+
+                allocate (t_array_${t1[0]}$${k1}$_${rank}$ :: wrapper%array, stat=stat)
+                if (stat /= 0) then
+                    msg = 'Failed to allocate array.'; return
+                end if
+
+                select type (typed_array => wrapper%array)
+                class is (t_array_${t1[0]}$${k1}$_${rank}$)
+                    typed_array%values = array
+                class default
+                    msg = 'Failed to allocate values.'; stat = 1; return
+                end select
+            end
+        #:endfor
+    #:endfor
+
     !> Version: experimental
     !>
     !> Return the positions of the true elements in array.
@@ -45,7 +79,7 @@ contains
         integer :: loc(count(array))
 
         call logicalloc(loc, array, .true., lbound)
-    end function trueloc
+    end
 
     !> Version: experimental
     !>
diff --git a/src/stdlib_io_np.fypp b/src/stdlib_io_np.fypp
@@ -70,7 +70,7 @@
 !> utf8-encoded string, so supports structured types with any unicode field names.
 module stdlib_io_np
     use stdlib_kinds, only: int8, int16, int32, int64, sp, dp, xdp, qp
-    use stdlib_array, only: t_array_bundle
+    use stdlib_array, only: t_array_wrapper
     implicit none
     private
 
@@ -121,9 +121,9 @@ module stdlib_io_np
     !> Load multiple multidimensional arrays from a (compressed) npz file.
     !> ([Specification](../page/specs/stdlib_io.html#load_npz))
     interface load_npz
-        module subroutine load_npz_to_bundle(filename, array_bundle, iostat, iomsg)
+        module subroutine load_npz_to_arrays(filename, arrays, iostat, iomsg)
             character(len=*), intent(in) :: filename
-            type(t_array_bundle), intent(out) :: array_bundle
+            type(t_array_wrapper), allocatable, intent(out) :: arrays(:)
             integer, intent(out), optional :: iostat
             character(len=:), allocatable, intent(out), optional :: iomsg
         end
@@ -134,9 +134,9 @@ module stdlib_io_np
     !> Save multidimensional arrays to a compressed or an uncompressed npz file.
     !> ([Specification](../page/specs/stdlib_io.html#save_npz))
     interface save_npz
-        module subroutine save_npz_from_bundle(filename, array_bundle, compressed, iostat, iomsg)
+        module subroutine save_npz_from_arrays(filename, arrays, compressed, iostat, iomsg)
             character(len=*), intent(in) :: filename
-            type(t_array_bundle), intent(in) :: array_bundle
+            type(t_array_wrapper), intent(in) :: arrays(*)
             !> If true, the file is saved in compressed format. The default is false.
             logical, intent(in), optional :: compressed
             integer, intent(out), optional :: iostat
diff --git a/src/stdlib_io_np_load.fypp b/src/stdlib_io_np_load.fypp
@@ -36,11 +36,27 @@ contains
 
         open(newunit=io, file=filename, form="unformatted", access="stream", iostat=stat)
         catch: block
+            character(len=:), allocatable :: this_type
             integer, allocatable :: vshape(:)
 
-            call verify_npy_file(io, filename, vtype, vshape, rank, stat, msg)
+            call get_descriptor(io, filename, this_type, vshape, stat, msg)
             if (stat /= 0) exit catch
 
+            if (this_type /= vtype) then
+                stat = 1
+                msg = "File '"//filename//"' contains data of type '"//this_type//"', "//&
+                    & "but expected '"//vtype//"'"
+                exit catch
+            end if
+
+            if (size(vshape) /= rank) then
+                stat = 1
+                msg = "File '"//filename//"' contains data of rank "//&
+                    & to_string(size(vshape))//", but expected "//&
+                    & to_string(rank)
+                exit catch
+            end if
+
             call allocate_array(array, vshape, stat)
             if (stat /= 0) then
                 msg = "Failed to allocate array of type '"//vtype//"' "//&
@@ -67,44 +83,6 @@ contains
   #:endfor
 #:endfor
 
-    !> Verify header, type and rank of the npy file.
-    subroutine verify_npy_file(io, filename, vtype, vshape, rank, stat, msg)
-        !> Access unit to the npy file.
-        integer, intent(in) :: io
-        !> Name of the npy file to load from.
-        character(len=*), intent(in) :: filename
-        !> Type of the data stored, retrieved from field `descr`.
-        character(len=*), intent(in) :: vtype
-        !> Shape of the stored data, retrieved from field `shape`.
-        integer, allocatable, intent(out) :: vshape(:)
-        !> Expected rank of the data.
-        integer, intent(in) :: rank
-        !> Status of operation.
-        integer, intent(out) :: stat
-        !> Associated error message in case of non-zero status.
-        character(len=:), allocatable, intent(out) :: msg
-
-        character(len=:), allocatable :: this_type
-
-        call get_descriptor(io, filename, this_type, vshape, stat, msg)
-        if (stat /= 0) return
-
-        if (this_type /= vtype) then
-            stat = 1
-            msg = "File '"//filename//"' contains data of type '"//this_type//"', "//&
-                & "but expected '"//vtype//"'"
-            return
-        end if
-
-        if (size(vshape) /= rank) then
-            stat = 1
-            msg = "File '"//filename//"' contains data of rank "//&
-                & to_string(size(vshape))//", but expected "//&
-                & to_string(rank)
-            return
-        end if
-    end
-
     #:for k1, t1 in KINDS_TYPES
         #:for rank in RANKS
             module subroutine allocate_array_${t1[0]}$${k1}$_${rank}$(array, vshape, stat)
@@ -126,9 +104,9 @@ contains
     !>
     !> Load multidimensional arrays from a compressed or uncompressed npz file.
     !> ([Specification](../page/specs/stdlib_io.html#load_npz))
-    module subroutine load_npz_to_bundle(filename, array_bundle, iostat, iomsg)
+    module subroutine load_npz_to_arrays(filename, arrays, iostat, iomsg)
         character(len=*), intent(in) :: filename
-        type(t_array_bundle), intent(out) :: array_bundle
+        type(t_array_wrapper), allocatable, intent(out) :: arrays(:)
         integer, intent(out), optional :: iostat
         character(len=:), allocatable, intent(out), optional :: iomsg
 
@@ -138,9 +116,9 @@ contains
 
         call unzip(filename, unzipped_bundle, stat, msg)
         if (stat == 0) then
-            call load_raw_to_bundle(unzipped_bundle, array_bundle, stat, msg)
+            call load_unzipped_bundle_to_arrays(unzipped_bundle, arrays, stat, msg)
         else
-            call identify_problem(filename, stat, msg)
+            call identify_unzip_problem(filename, stat, msg)
         end if
 
         if (present(iostat)) then
@@ -156,64 +134,77 @@ contains
         if (present(iomsg) .and. allocated(msg)) call move_alloc(msg, iomsg)
     end
 
-    module subroutine load_raw_to_bundle(unzipped_bundle, array_bundle, stat, msg)
+    module subroutine load_unzipped_bundle_to_arrays(unzipped_bundle, arrays, stat, msg)
         type(t_unzipped_bundle), intent(in) :: unzipped_bundle
-        type(t_array_bundle), intent(out) :: array_bundle
+        type(t_array_wrapper), allocatable, intent(out) :: arrays(:)
         integer, intent(out) :: stat
         character(len=:), allocatable, intent(out) :: msg
 
         integer :: i, io
+        integer, allocatable :: vshape(:)
+        character(len=:), allocatable :: this_type
+
+        allocate (arrays(size(unzipped_bundle%files)))
 
-        allocate (array_bundle%files(size(unzipped_bundle%files)))
         do i = 1, size(unzipped_bundle%files)
-            array_bundle%files(i)%name = unzipped_bundle%files(i)%name
-            open (newunit=io, status='scratch', form='unformatted', access='stream', iostat=stat)
-            if (stat /= 0) return
-            write (io) unzipped_bundle%files(i)%data
-            call load_string_to_array(io, unzipped_bundle%files(i), array_bundle%files(i), stat, msg)
-            close (io, status='delete', iostat=stat)
+            open (newunit=io, status='scratch', form='unformatted', access='stream', iostat=stat, iomsg=msg)
             if (stat /= 0) return
-        end do
-    end
-
-    module subroutine load_string_to_array(io, unzipped_file, array, stat, msg)
-        integer, intent(in) :: io
-        type(t_unzipped_file), intent(in) :: unzipped_file
-        class(t_array), intent(inout) :: array
-        integer, intent(out) :: stat
-        character(len=:), allocatable, intent(out) :: msg
 
-        #:for k1, t1 in KINDS_TYPES
-            #:for rank in RANKS
-                ${t1}$, allocatable :: array_${t1[0]}$${k1}$_${rank}$${ranksuffix(rank)}$
-            #:endfor
-        #:endfor
+            write (io, iostat=stat) unzipped_bundle%files(i)%data
+            if (stat /= 0) then
+                msg = 'Failed to write unzipped data to scratch file.'
+                close (io, status='delete'); return
+            end if
 
-        integer, allocatable :: vshape(:)
+            rewind (io)
+            call get_descriptor(io, unzipped_bundle%files(i)%name, this_type, vshape, stat, msg)
+            if (stat /= 0) return
 
-        select type (arr => array)
+            select case (this_type)
             #:for k1, t1 in KINDS_TYPES
+            case (type_${t1[0]}$${k1}$)
+                select case (size(vshape))
                 #:for rank in RANKS
-                    type is (t_array_${t1[0]}$${k1}$_${rank}$)
-                    call verify_npy_file(io, unzipped_file%name, type_${t1[0]}$${k1}$, vshape, ${rank}$, stat, msg)
-                    if (stat /= 0) return
-                    call allocate_array(array_${t1[0]}$${k1}$_${rank}$, vshape, stat)
-                    if (stat /= 0) then
-                        msg = "Failed to allocate array of type '"//type_${t1[0]}$${k1}$//"' "//&
-                            & "with total size of "//to_string(product(vshape))
-                        return
-                    end if
-                    read (io, iostat=stat) array_${t1[0]}$${k1}$_${rank}$${ranksuffix(rank)}$
-                    arr%values = array_${t1[0]}$${k1}$_${rank}$${ranksuffix(rank)}$
+                case (${rank}$)
+                    block
+                        ${t1}$, allocatable :: array${ranksuffix(rank)}$
+
+                        call allocate_array(array, vshape, stat)
+                        if (stat /= 0) then
+                            msg = "Failed to allocate array of type '"//this_type//"'."; return
+                        end if
+
+                        read (io, iostat=stat) array
+                        if (stat /= 0) then
+                            msg = "Failed to read array of type '"//this_type//"' "//&
+                                & 'with total size of '//to_string(product(vshape)); return
+                        end if
+
+                        call arrays(i)%allocate_array(array, stat, msg)
+                        if (stat /= 0) then
+                            msg = "Failed to allocate array of type '"//this_type//"' "//&
+                                & 'with total size of '//to_string(product(vshape)); return
+                        end if
+
+                        arrays(i)%array%name = unzipped_bundle%files(i)%name
+                    end block
                 #:endfor
+                case default
+                    stat = 1; msg = 'Unsupported rank for array of type '//this_type//': '// &
+                    & to_string(size(vshape))//'.'; return
+                end select
             #:endfor
-        class default
-            stat = 1; msg = 'Unsupported array type.'; return
-        end select
+            case default
+                stat = 1; msg = 'Unsupported array type: '//this_type//'.'; return
+            end select
+
+            close (io, status='delete')
+            if (stat /= 0) return
+        end do
     end
 
-    !> Open file and try to identify the problem.
-    module subroutine identify_problem(filename, stat, msg)
+    !> Open file and try to identify the cause of the error that occurred during unzip.
+    module subroutine identify_unzip_problem(filename, stat, msg)
         character(len=*), intent(in) :: filename
         integer, intent(inout) :: stat
         character(len=:), allocatable, intent(inout) :: msg
@@ -291,7 +282,7 @@ contains
         
         ! stat should be zero if no error occurred
         stat = 0
-        
+
         read(io, iostat=stat) header
         if (stat /= 0) return
 
diff --git a/src/stdlib_io_np_save.fypp b/src/stdlib_io_np_save.fypp
@@ -135,9 +135,9 @@ contains
     !>
     !> Save multidimensional arrays to a compressed or an uncompressed npz file.
     !> ([Specification](../page/specs/stdlib_io.html#save_npz))
-    module subroutine save_npz_from_bundle(filename, array_bundle, compressed, iostat, iomsg)
+    module subroutine save_npz_from_arrays(filename, arrays, compressed, iostat, iomsg)
         character(len=*), intent(in) :: filename
-        type(t_array_bundle), intent(in) :: array_bundle
+        type(t_array_wrapper), intent(in) :: arrays(*)
         !> If true, the file is saved in compressed format. The default is false.
         logical, intent(in), optional :: compressed
         integer, intent(out), optional :: iostat
diff --git a/src/stdlib_io_zip.f90 b/src/stdlib_io_zip.f90
@@ -1,5 +1,4 @@
 module stdlib_io_zip
-    use stdlib_array, only: t_array_bundle
     use stdlib_io_minizip
     use iso_c_binding, only: c_ptr, c_associated, c_int, c_long, c_char
     implicit none
@@ -13,7 +12,7 @@ module stdlib_io_zip
     integer(kind=c_long), parameter :: buffer_size = 1024
 
     interface unzip
-        module procedure unzip_to_raw
+        module procedure unzip_to_bundle
     end interface
 
     !> Contains extracted raw data from a zip file.
@@ -32,7 +31,7 @@ module stdlib_io_zip
 
 contains
 
-    module subroutine unzip_to_raw(filename, bundle, iostat, iomsg)
+    module subroutine unzip_to_bundle(filename, bundle, iostat, iomsg)
         character(len=*), intent(in) :: filename
         type(t_unzipped_bundle), intent(out) :: bundle
         integer, intent(out), optional :: iostat
diff --git a/test/io/test_np.f90 b/test/io/test_np.f90
