From 6ed8b5c479000cd76994f9f45a70f2f6128d6550 Mon Sep 17 00:00:00 2001
From: Brett <brettgraham@gmail.com>
Date: Fri, 11 Oct 2024 11:09:43 -0400
Subject: [PATCH] add optional out to read_into_array

---
 asdf/_block/io.py  |  3 +--
 asdf/generic_io.py | 28 +++++++++++++++++++++-------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/asdf/_block/io.py b/asdf/_block/io.py
index e38c4377d..5ccd4d724 100644
--- a/asdf/_block/io.py
+++ b/asdf/_block/io.py
@@ -158,8 +158,7 @@ def read_block_data(fd, header, offset=None, memmap=False, out=None):
             data = fd.memmap_array(offset, used_size)
             ff_bytes = header["allocated_size"]
         else:
-            # TODO update to read into out
-            data = fd.read_into_array(used_size)
+            data = fd.read_into_array(used_size, out=out)
             ff_bytes = header["allocated_size"] - header["used_size"]
         if (header["flags"] & constants.BLOCK_FLAG_STREAMED) and fd.seekable():
             fd.seek(0, os.SEEK_END)
diff --git a/asdf/generic_io.py b/asdf/generic_io.py
index 4114e104f..00545d0f9 100644
--- a/asdf/generic_io.py
+++ b/asdf/generic_io.py
@@ -669,7 +669,7 @@ def flush_memmap(self):
         msg = f"memmapping is not implemented for {self.__class__.__name__}"
         raise NotImplementedError(msg)
 
-    def read_into_array(self, size):
+    def read_into_array(self, size, out=None):
         """
         Read a chunk of the file into a uint8 array.
 
@@ -683,7 +683,10 @@ def read_into_array(self, size):
         array : np.memmap
         """
         buff = self.read(size)
-        return np.frombuffer(buff, np.uint8, size, 0)
+        if out is None:
+            return np.frombuffer(buff, np.uint8, size, 0)
+        out[:] = np.frombuffer(buff, np.uint8, size, 0)
+        return out
 
 
 class GenericWrapper:
@@ -794,7 +797,13 @@ def flush_memmap(self):
         if hasattr(self, "_mmap"):
             self._mmap.flush()
 
-    def read_into_array(self, size):
+    def read_into_array(self, size, out=None):
+        if out is not None:
+            read_size = self._fd.readinto(out.data)
+            if read_size != size:
+                # TODO better message here
+                raise OSError("Read inconsistency")
+            return out
         return np.fromfile(self._fd, dtype=np.uint8, count=size)
 
     def _fix_permissions(self):
@@ -853,13 +862,17 @@ class MemoryIO(RandomAccessFile):
     def __init__(self, fd, mode, uri=None):
         super().__init__(fd, mode, uri=uri)
 
-    def read_into_array(self, size):
+    def read_into_array(self, size, out=None):
         buf = self._fd.getvalue()
         offset = self._fd.tell()
+        # TODO improve this
         result = np.frombuffer(buf, np.uint8, size, offset)
-        # Copy the buffer so the original memory can be released.
-        result = result.copy()
         self.seek(size, SEEK_CUR)
+        if out is not None:
+            out[:] = result
+        else:
+            # Copy the buffer so the original memory can be released.
+            result = result.copy()
         return result
 
 
@@ -937,7 +950,8 @@ def fast_forward(self, size):
             msg = "Read past end of file"
             raise OSError(msg)
 
-    def read_into_array(self, size):
+    def read_into_array(self, size, out=None):
+        # TODO support out... what becomes an InputStream?
         try:
             # See if Numpy can handle this as a real file first...
             return np.fromfile(self._fd, np.uint8, size)