Respond to review

Author: oerc0122

castep_outputs/__init__.py

@@ -7,3 +7,5 @@
 
 from .cli.castep_outputs_main import parse_single  # noqa: F401
 from .parsers import *  # noqa: F403
+from .tools import MDGeomParser as MDGeomParser  # noqa: PLC0414
+from .tools import get_generated_files as get_generated_files  # noqa: PLC0414

castep_outputs/tools/__init__.py

@@ -2,3 +2,7 @@
 
 For more advanced scripts, see ``castep_outputs_tools``.
 """
+
+
+from .get_generated_files import get_generated_files as get_generated_files  # noqa: PLC0414
+from .md_geom_parser import MDGeomParser as MDGeomParser  # noqa: PLC0414

castep_outputs/tools/md_geom_parser.py

@@ -5,7 +5,7 @@
 from collections.abc import Generator, Sequence
 from functools import singledispatchmethod
 from pathlib import Path
-from typing import Union, overload
+from typing import overload
 
 from castep_outputs.parsers.md_geom_file_parser import (
     MDGeomTimestepInfo,
@@ -29,7 +29,7 @@ class MDGeomParser:
     """
 
     def __init__(self, md_geom_file: Path | str) -> None:
-        self._next_frame = 0
+        self._next_frame: int | None
 
         self.file = Path(md_geom_file).expanduser()
 
@@ -52,12 +52,12 @@ def __init__(self, md_geom_file: Path | str) -> None:
 
         while next(self._handle).strip():
             pass
-        self._frame_len = self._handle.lineno - self._start_line - 1
+        self._frame_lines = self._handle.lineno - self._start_line - 1
 
-        self._byte_len = self._handle.tell() - self._start
+        self._frame_bytes = self._handle.tell() - self._start
         stat = self.file.stat()
 
-        len_est = (stat.st_size - self._start) / self._byte_len
+        len_est = (stat.st_size - self._start) / self._frame_bytes
 
         if not len_est.is_integer():
             self.logger(
@@ -71,10 +71,11 @@ def __init__(self, md_geom_file: Path | str) -> None:
             )
 
         self._len = int(len_est)
+        self._go_to_frame(0)
 
     @property
-    def next_frame(self) -> int:
-        """Get index of next frame to be read."""
+    def next_frame(self) -> int | None:
+        """Get index of next frame to be read, or None if at file end."""
         return self._next_frame
 
     def _get_index(self, frame: int) -> int:
@@ -90,13 +91,13 @@ def _get_index(self, frame: int) -> int:
         int
             Current index in bytes.
         """
-        return self._start + (self._byte_len * frame)
+        return self._start + (self._frame_bytes * frame)
 
     def _go_to_frame(self, frame: int) -> None:
         """Set file pointer to given index."""
         ind = self._get_index(frame)
         self._handle.file.seek(ind)
-        self._handle._lineno = self._start_line + (frame * self._frame_len)
+        self._handle._lineno = self._start_line + (frame * self._frame_lines)
         self._next_frame = frame if frame < len(self) else None
 
     def get_frame(self, frame: int) -> MDGeomTimestepInfo:
@@ -111,16 +112,20 @@ def get_frame(self, frame: int) -> MDGeomTimestepInfo:
         -------
         MDGeomTimestepInfo
             Parsed frame.
+
+        Raises
+        ------
+        IndexError
+            Requested frame out of range.
         """
-        if -len(self) > frame > len(self):
-            self.logger(
-                f"Cannot get {frame}th frame. File only has {len(self)} frames.", level="warning",
-            )
+        if -len(self) > frame < len(self):
+            raise IndexError(f"Cannot get {frame}th frame. File only has {len(self)} frames.")
 
         if frame < 0:
             frame = len(self) + frame
 
-        self._go_to_frame(frame)
+        if frame != self.next_frame:
+            self._go_to_frame(frame)
 
         return next(self)
 
@@ -145,13 +150,12 @@ def __iter__(self) -> Generator[MDGeomTimestepInfo, int, None]:
         MDGeomTimestepInfo
             Frames in file.
         """
-        self._handle.file.seek(self._start)
-        self._handle._lineno = self._start_line
-        self._next_frame = 0
-        while self._next_frame is not None:
-            jump = yield next(self)
-            if jump is not None:
-                self._go_to_frame(jump)
+        i = 0
+        while i < len(self):
+            trial = yield self[i]
+            i += 1
+            if trial is not None:
+                i = trial
 
     def __next__(self) -> MDGeomTimestepInfo:
         """Get the next frame.
@@ -166,7 +170,7 @@ def __next__(self) -> MDGeomTimestepInfo:
         StopIteration
             No next frame.
         """
-        if not (block := Block.get_lines(self._handle, self._frame_len, eof_possible=True)):
+        if not (block := Block.get_lines(self._handle, self._frame_lines, eof_possible=True)):
             raise StopIteration
 
         self._next_frame = self._next_frame + 1 if self._next_frame < len(self) - 1 else None
@@ -176,7 +180,7 @@ def __next__(self) -> MDGeomTimestepInfo:
     def __getitem__(self, frame: int) -> MDGeomTimestepInfo: ...
 
     @overload
-    def __getitem__(self, frame: Union[Sequence, slice]) -> list[MDGeomTimestepInfo]: ...
+    def __getitem__(self, frame: Sequence | slice) -> list[MDGeomTimestepInfo]: ...
 
     @singledispatchmethod
     def __getitem__(self, frame) -> list[MDGeomTimestepInfo] | MDGeomTimestepInfo:

pyproject.toml

@@ -109,7 +109,6 @@ ignore = [
     "PLR2004",  # Magic constant value
     "D105",     # Magic method doc
     "ANN401",   # No typing.Any
-    "UP007",    # Use | union
 ]
 allowed-confusables = ["α", "β", "γ"]