refactoring of calculate_number_of_frames

PQAnalysis/io/traj_file/trajectory_reader.py

@@ -6,7 +6,6 @@
 
 # 3rd party modules
 import logging
-import os
 from beartype.typing import List, Generator
 from tqdm.auto import tqdm
 
@@ -238,7 +237,8 @@ def frame_generator(
                         frame_lines = [line]
 
                 if frame_lines:
-                    frame = self._read_single_frame("".join(frame_lines), self.topology)
+                    frame = self._read_single_frame(
+                        "".join(frame_lines), self.topology)
 
                     if frame.cell.is_vacuum and last_cell is not None:
                         frame.cell = last_cell
@@ -398,22 +398,44 @@ def calculate_number_of_frames(self) -> int:
         n_frames = 0
 
         for filename in self.filenames:
-
-            # TODO: Add check to BaseReader to check if the file is empty
-            if os.path.getsize(filename) == 0:
-                continue
-
             with open(filename, "r", encoding="utf-8") as f:
 
-                # Read the lines
                 lines = f.readlines()
-
-            # +2 for the cell and atom count lines
-            n_frames += int(len(lines) / (int(lines[0].split()[0]) + 2))
+                n_lines = len(lines)
+
+                # If the file is empty, continue to the next file
+                if n_lines == 0:
+                    continue
+
+                try:
+                    n_atoms = int(lines[0].split()[0])
+                except (ValueError, IndexError):
+                    self.logger.error(
+                        (
+                            "Invalid number of atoms in the first line "
+                            f"of file {filename}."
+                        ),
+                        exception=TrajectoryReaderError
+                    )
+
+                # +2 for the cell and atom count lines
+                _n_frames, remainder = divmod(n_lines, n_atoms + 2)
+
+                if remainder != 0:
+                    self.logger.error(
+                        (
+                            "The number of lines in the file is not divisible "
+                            f"by the number of atoms {n_atoms} "
+                            "in the first line."
+                        ),
+                        exception=TrajectoryReaderError
+                    )
+
+                n_frames += _n_frames
 
         return n_frames
 
-    @property
+    @ property
     def cells(self) -> list[Cell]:
         """
         Returns the cells of the trajectory.

tests/io/test_trajectoryReader.py

@@ -156,10 +156,12 @@ def test_frame_generator(self, caplog):
         test_frames = [frame for frame in reader.frame_generator()]
         assert test_frames == [frame1, frame2, frame1, frame2]
 
-        test_frames = [frame for frame in reader.frame_generator(trajectory_start=1)]
+        test_frames = [
+            frame for frame in reader.frame_generator(trajectory_start=1)]
         assert test_frames == [frame2, frame1, frame2]
 
-        test_frames = [frame for frame in reader.frame_generator(trajectory_stop=3)]
+        test_frames = [
+            frame for frame in reader.frame_generator(trajectory_stop=3)]
         assert test_frames == [frame1, frame2, frame1]
 
         # TODO: test topology set when const_topology and topology is None (Pylint)
@@ -191,7 +193,8 @@ def test_frame_generator(self, caplog):
             TrajectoryReader.__qualname__,
             exception=IndexError,
             logging_level="ERROR",
-            message_to_test=("start index is greater than or equal to the stop index"),
+            message_to_test=(
+                "start index is greater than or equal to the stop index"),
             function=reader.frame_generator(
                 trajectory_start=1, trajectory_stop=1
             ).__next__,
@@ -245,7 +248,8 @@ def test_window_generator(self, caplog):
         test_frames = list(reader.window_generator(window_size=4))
         assert test_frames == [Trajectory([frame1, frame2, frame1, frame2])]
 
-        test_frames = list(reader.window_generator(window_size=2, window_gap=2))
+        test_frames = list(reader.window_generator(
+            window_size=2, window_gap=2))
         assert test_frames == [
             Trajectory([frame1, frame2]),
             Trajectory([frame1, frame2]),
@@ -318,7 +322,8 @@ def test_window_generator(self, caplog):
             message_to_test=(
                 "start index is less than 0 or greater than the length of the trajectory"
             ),
-            function=reader.window_generator(1, 1, trajectory_start=-1).__next__,
+            function=reader.window_generator(
+                1, 1, trajectory_start=-1).__next__,
         )
 
         assert_logging_with_exception(
@@ -342,7 +347,8 @@ def test_window_generator(self, caplog):
             message_to_test=(
                 "stop index is less than 0 or greater than the length of the trajectory"
             ),
-            function=reader.window_generator(1, 1, trajectory_stop=-1).__next__,
+            function=reader.window_generator(
+                1, 1, trajectory_stop=-1).__next__,
         )
 
         assert_logging_with_exception(
@@ -363,7 +369,8 @@ def test_window_generator(self, caplog):
             TrajectoryReader.__qualname__,
             exception=IndexError,
             logging_level="ERROR",
-            message_to_test=("start index is greater than or equal to the stop index"),
+            message_to_test=(
+                "start index is greater than or equal to the stop index"),
             function=reader.window_generator(
                 1, 1, trajectory_start=1, trajectory_stop=0
             ).__next__,
@@ -374,7 +381,8 @@ def test_window_generator(self, caplog):
             TrajectoryReader.__qualname__,
             exception=IndexError,
             logging_level="ERROR",
-            message_to_test=("start index is greater than or equal to the stop index"),
+            message_to_test=(
+                "start index is greater than or equal to the stop index"),
             function=reader.window_generator(
                 1, 1, trajectory_start=1, trajectory_stop=1
             ).__next__,

tests/traj/test_trajectory.py

@@ -238,7 +238,8 @@ def test_window(self, caplog):
             logging_level="ERROR",
             message_to_test=(
                 "start index is greater than or equal to the stop index"),
-            function=traj.window(1, 1, trajectory_start=2, trajectory_stop=1).__next__,
+            function=traj.window(1, 1, trajectory_start=2,
+                                 trajectory_stop=1).__next__,
         )
 
         assert_logging_with_exception(
@@ -249,7 +250,8 @@ def test_window(self, caplog):
             message_to_test=(
                 "window size is greater than the trajectory_stop - trajectory_start"
             ),
-            function=traj.window(3, 1, trajectory_start=1, trajectory_stop=3).__next__,
+            function=traj.window(3, 1, trajectory_start=1,
+                                 trajectory_stop=3).__next__,
         )
 
     def test__iter__(self):
@@ -298,7 +300,7 @@ def test_append(self):
 
         assert traj.frames == [self.frame1]
 
-    def test_pop(self): 
+    def test_pop(self):
         traj = Trajectory(self.frames)
         assert traj.pop() == self.frame3
         assert traj.frames == [self.frame1, self.frame2]