From 3f2da1bba9079b3287a99dff0de919fce26349a2 Mon Sep 17 00:00:00 2001
From: nain <126972030+F49FF806@users.noreply.github.com>
Date: Fri, 26 May 2023 02:33:31 -0400
Subject: [PATCH] Simplify generator functions and make get_segment_dirs inputs
 consistent with get_segment_files

+ get_segment_dirs explicitly takes data_dir for clarity and future flexibility
+ removed multiple cases for one general purpose generator expression
---
 src/borg/repository.py | 42 +++++++++++++++++++-----------------------
 1 file changed, 19 insertions(+), 23 deletions(-)

diff --git a/src/borg/repository.py b/src/borg/repository.py
index 203596c85..2910ce261 100644
--- a/src/borg/repository.py
+++ b/src/borg/repository.py
@@ -1319,41 +1319,37 @@ class LoggedIO:
         safe_fadvise(fd.fileno(), 0, 0, 'DONTNEED')
         fd.close()
 
-    def get_segment_dirs(self, start_index=None, end_index=None):
-        """Returns generator yeilding required segment dirs as `os.DirEntry` objects. Start and end are inclusive.
+    def get_segment_dirs(self, data_dir, start_index=0, end_index=2**30):
+        """Returns generator yielding required segment dirs in data_dir as `os.DirEntry` objects.
+        Start and end are inclusive.
         """
-        data_dir = os.path.join(self.path, 'data')
-        if start_index is None and end_index is None:
-            segment_dirs = (f for f in os.scandir(data_dir) if f.is_dir() and f.name.isdigit())
-        elif start_index is not None and end_index is None:
-            segment_dirs = (f for f in os.scandir(data_dir) if f.is_dir() and f.name.isdigit() and start_index <= int(f.name))
-        elif start_index is None and end_index is not None:
-            segment_dirs = (f for f in os.scandir(data_dir) if f.is_dir() and f.name.isdigit() and int(f.name) <= end_index)
-        elif start_index is not None and end_index is not None:
-            segment_dirs = (f for f in os.scandir(data_dir) if f.is_dir() and f.name.isdigit() and start_index <= int(f.name) <= end_index)
+        segment_dirs = (
+            f
+            for f in os.scandir(data_dir)
+            if f.is_dir() and f.name.isdigit() and start_index <= int(f.name) <= end_index
+        )
         return segment_dirs
 
-    def get_segment_files(self, segment_dir, start_index=None, end_index=None):
-        """Returns generator yeilding required segment files in segment_dir as `os.DirEntry` objects. Start and end are inclusive.
+    def get_segment_files(self, segment_dir, start_index=0, end_index=2**32):
+        """Returns generator yielding required segment files in segment_dir as `os.DirEntry` objects.
+        Start and end are inclusive.
         """
-        if start_index is None and end_index is None:
-            segment_files = (f for f in os.scandir(segment_dir) if f.is_file() and f.name.isdigit())
-        elif start_index is not None and end_index is None:
-            segment_files = (f for f in os.scandir(segment_dir) if f.is_file() and f.name.isdigit() and start_index <= int(f.name))
-        elif start_index is None and end_index is not None:
-            segment_files = (f for f in os.scandir(segment_dir) if f.is_file() and f.name.isdigit() and int(f.name) <= end_index)
-        elif start_index is not None and end_index is not None:
-            segment_files = (f for f in os.scandir(segment_dir) if f.is_file() and f.name.isdigit() and start_index <= int(f.name) <= end_index)
+        segment_files = (
+            f
+            for f in os.scandir(segment_dir)
+            if f.is_file() and f.name.isdigit() and start_index <= int(f.name) <= end_index
+        )
         return segment_files
 
     def segment_iterator(self, segment=None, reverse=False):
         if segment is None:
             segment = 0 if not reverse else 2 ** 32 - 1
         start_segment_dir = segment // self.segments_per_dir
+        data_path = os.path.join(self.path, 'data')
         if not reverse:
-            dirs = self.get_segment_dirs(start_index=start_segment_dir)
+            dirs = self.get_segment_dirs(data_path, start_index=start_segment_dir)
         else:
-            dirs = self.get_segment_dirs(end_index=start_segment_dir)
+            dirs = self.get_segment_dirs(data_path, end_index=start_segment_dir)
         dirs = sorted(dirs, key=lambda dir: int(dir.name), reverse=reverse)
         for dir in dirs:
             if not reverse: