read_n returns remaining frames (#105)

Modified read_n to return the number of frames available if less than
the number of frames requested.

```python
#f is a CtbRawFile containing 10 frames

f.read_n(7) # you get 7 frames
f.read_n(7) # you get 3 frames
f.read_n(7) # RuntimeError
```

Also added support for chunk_size when iterating over a file:

```python
# The file contains 10 frames


with CtbRawFile(fname, chunk_size = 7) as f:
    for headers, frames in f:
        #do something with the data
        # 1 iteration 7 frames
        # 2 iteration 3 frames
        # 3 iteration stops
```
This commit is contained in:
heymes_j 2024-11-26 14:07:21 +01:00 committed by GitHub
commit 06670a7e24
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 4 deletions

View File

@ -2,6 +2,7 @@
from . import _aare from . import _aare
import numpy as np import numpy as np
from .ScanParameters import ScanParameters from .ScanParameters import ScanParameters
class CtbRawFile(_aare.CtbRawFile): class CtbRawFile(_aare.CtbRawFile):
"""File reader for the CTB raw file format. """File reader for the CTB raw file format.
@ -11,9 +12,10 @@ class CtbRawFile(_aare.CtbRawFile):
The function should take a numpy array of type uint8 and return one The function should take a numpy array of type uint8 and return one
or several numpy arrays. or several numpy arrays.
""" """
def __init__(self, fname, transform = None): def __init__(self, fname, transform = None, chunk_size = 1):
super().__init__(fname) super().__init__(fname)
self.transform = transform self.transform = transform
self._chunk_size = chunk_size
def read_frame(self, frame_index: int | None = None ) -> tuple: def read_frame(self, frame_index: int | None = None ) -> tuple:
@ -59,6 +61,10 @@ class CtbRawFile(_aare.CtbRawFile):
Uses the position of the file pointer :py:meth:`~CtbRawFile.tell` to determine Uses the position of the file pointer :py:meth:`~CtbRawFile.tell` to determine
where to start reading from. where to start reading from.
If the number of frames requested is larger than the number of frames left in the file,
the function will read the remaining frames. If no frames are left in the file
a RuntimeError is raised.
Args: Args:
n_frames (int): Number of frames to read. n_frames (int): Number of frames to read.
@ -68,6 +74,12 @@ class CtbRawFile(_aare.CtbRawFile):
Raises: Raises:
RuntimeError: If EOF is reached. RuntimeError: If EOF is reached.
""" """
# Calculate the number of frames to actually read
n_frames = min(n_frames, self.frames_in_file - self.tell())
if n_frames == 0:
raise RuntimeError("No frames left in file.")
# Do the first read to figure out what we have # Do the first read to figure out what we have
tmp_header, tmp_data = self.read_frame() tmp_header, tmp_data = self.read_frame()
@ -87,10 +99,12 @@ class CtbRawFile(_aare.CtbRawFile):
def read(self) -> tuple: def read(self) -> tuple:
"""Read the entire file. """Read the entire file.
Seeks to the beginning of the file before reading.
Returns: Returns:
tuple: header, data tuple: header, data
""" """
self.seek(0)
return self.read_n(self.frames_in_file) return self.read_n(self.frames_in_file)
def seek(self, frame_index:int) -> None: def seek(self, frame_index:int) -> None:
@ -101,7 +115,7 @@ class CtbRawFile(_aare.CtbRawFile):
""" """
super().seek(frame_index) super().seek(frame_index)
def tell() -> int: def tell(self) -> int:
"""Return the current frame position in the file. """Return the current frame position in the file.
Returns: Returns:
@ -164,7 +178,12 @@ class CtbRawFile(_aare.CtbRawFile):
def __next__(self): def __next__(self):
try: try:
if self._chunk_size == 1:
return self.read_frame() return self.read_frame()
else:
return self.read_n(self._chunk_size)
except RuntimeError: except RuntimeError:
# TODO! find a good way to check that we actually have the right exception # TODO! find a good way to check that we actually have the right exception
raise StopIteration raise StopIteration