Remove duplicated file queries in pipeline staging

src/tigerflow/pipeline.py

@@ -249,43 +249,38 @@ def _start_tasks(self):
             else:
                 raise ValueError(f"Unsupported task kind: {type(task)}")
 
-    def _build_staging_context(self) -> StagingContext:
-        """Build the current context for staging middleware."""
-        failed_stems = self._failed_stems()
+    def _prepare_staging_inputs(self) -> tuple[list[Path], StagingContext]:
+        """Collect the candidates and context for staging middleware.
 
-        n_finished = sum(1 for f in self._finished_dir.iterdir() if f.is_file())
+        Both are built here so the input directory is scanned once: `waiting`
+        is derived from the same list the middleware chain receives.
+        """
+        failed_stems = self._failed_stems()
         n_staged = sum(
             1
             for f in self._symlinks_dir.iterdir()
             if f.is_file()
             and f.name.removesuffix(self._config.root_input_ext) not in failed_stems
         )
-        n_waiting = sum(
-            1
+        candidates = [
+            f
             for f in self._input_dir.iterdir()
             if f.is_file()
             and f.name.endswith(self._config.root_input_ext)
             and f.name not in self._filenames
-        )
-
-        return StagingContext(
-            waiting=n_waiting,
+        ]
+        context = StagingContext(
+            waiting=len(candidates),
             staged=n_staged,
-            completed=n_finished,
+            completed=self._count_finished(),
             failed=len(failed_stems),
             input_dir=self._input_dir,
             output_dir=self._output_dir,
         )
+        return candidates, context
 
     def _stage_new_files(self):
-        context = self._build_staging_context()
-        candidates = [
-            f
-            for f in self._input_dir.iterdir()
-            if f.is_file()
-            and f.name.endswith(self._config.root_input_ext)
-            and f.name not in self._filenames
-        ]
+        candidates, context = self._prepare_staging_inputs()
         to_stage = self._config.staging.process(candidates, context)
         for file in to_stage:
             self._symlinks_dir.joinpath(file.name).symlink_to(file)
@@ -406,6 +401,19 @@ def _report_failed_files(self):
             if n_files > 0:
                 logger.error("[{}] {} failed files", task.name, n_files)
 
+    def _count_finished(self) -> int:
+        return sum(1 for file in self._finished_dir.iterdir() if file.is_file())
+
+    def _all_tracked_files_settled(self) -> bool:
+        """Whether every tracked file has either finished or failed.
+
+        `_filenames` is never pruned, so it covers every file the pipeline has
+        ever tracked.
+        """
+        return self._count_finished() + len(self._failed_stems()) >= len(
+            self._filenames
+        )
+
     def _failed_stems(self) -> set[str]:
         """Stems of input files that failed in at least one task.
 
@@ -465,13 +473,11 @@ def _handle_processed_files(self):
         # Log progress
         if completed_file_ids:
             logger.info("Completed processing {} files", len(completed_file_ids))
-            n_finished = sum(1 for f in self._finished_dir.iterdir() if f.is_file())
-            if (n_finished + len(self._failed_stems())) >= len(self._filenames):
+            if self._all_tracked_files_settled():
                 logger.info("No more files to process, starting idle time count")
 
     def _check_inactivity(self):
-        n_finished = sum(1 for file in self._finished_dir.iterdir() if file.is_file())
-        if (n_finished + len(self._failed_stems())) < len(self._filenames):
+        if not self._all_tracked_files_settled():
             self._last_active = datetime.now()
 
         inactivity = datetime.now() - self._last_active

tests/integration/pipeline/test_file_staging.py

@@ -143,7 +143,7 @@ def test_counts_waiting_files(
         for i in range(3):
             (input_dir / f"f{i}.txt").write_text("x")
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
 
         assert context.waiting == 3
         assert context.staged == 0
@@ -157,7 +157,7 @@ def test_counts_shift_after_staging(
         (input_dir / "a.txt").write_text("x")
 
         pipeline._stage_new_files()
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
 
         assert context.waiting == 0
         assert context.staged == 1
@@ -174,7 +174,7 @@ def test_counts_completed_files(
         (task.output_dir / "a.txt").write_text("done")
         pipeline._handle_processed_files()
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
 
         assert context.completed == 1
         assert context.staged == 0
@@ -191,7 +191,7 @@ def test_staged_count_excludes_failures(
         (task.output_dir / "a.err").write_text("boom")
         pipeline._report_failed_files()
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
 
         assert context.failed == 1
         assert context.staged == 0
@@ -240,7 +240,7 @@ def test_counts_stay_consistent_in_mixed_state(
         pipeline._report_failed_files()
         pipeline._handle_processed_files()
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
 
         counts = (context.waiting, context.staged, context.completed, context.failed)
         assert counts == (0, 2, 1, 1)
@@ -286,7 +286,7 @@ def test_fan_out_failure_counts_one_file(
             "Failure must not remove the symlink the remaining task reads from"
         )
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
         assert (context.staged, context.failed) == (1, 1)
 
     def test_failed_file_excluded_under_multi_part_extension(
@@ -312,7 +312,7 @@ def test_failed_file_excluded_under_multi_part_extension(
         (task.output_dir / "sample.err").write_text("boom")
         pipeline._report_failed_files()
 
-        context = pipeline._build_staging_context()
+        _, context = pipeline._prepare_staging_inputs()
         assert (context.staged, context.failed) == (0, 1)