Skip to content

Commit

Permalink
add optional out to read_into_array
Browse files Browse the repository at this point in the history
  • Loading branch information
braingram committed Oct 11, 2024
1 parent 63a8c89 commit 6ed8b5c
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 9 deletions.
3 changes: 1 addition & 2 deletions asdf/_block/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
28 changes: 21 additions & 7 deletions asdf/generic_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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

Check warning on line 689 in asdf/generic_io.py

View check run for this annotation

Codecov / codecov/patch

asdf/generic_io.py#L688-L689

Added lines #L688 - L689 were not covered by tests


class GenericWrapper:
Expand Down Expand Up @@ -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:

Check warning on line 803 in asdf/generic_io.py

View check run for this annotation

Codecov / codecov/patch

asdf/generic_io.py#L802-L803

Added lines #L802 - L803 were not covered by tests
# TODO better message here
raise OSError("Read inconsistency")
return out

Check warning on line 806 in asdf/generic_io.py

View check run for this annotation

Codecov / codecov/patch

asdf/generic_io.py#L805-L806

Added lines #L805 - L806 were not covered by tests
return np.fromfile(self._fd, dtype=np.uint8, count=size)

def _fix_permissions(self):
Expand Down Expand Up @@ -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

Check warning on line 872 in asdf/generic_io.py

View check run for this annotation

Codecov / codecov/patch

asdf/generic_io.py#L872

Added line #L872 was not covered by tests
else:
# Copy the buffer so the original memory can be released.
result = result.copy()
return result


Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 6ed8b5c

Please sign in to comment.