changeset 52923:307c4a0b91a0

stream-clone-v2: turn the file chunking function into a class We did this in two step: extracting the code to a function, then turning the function into a class for clarity. The introduction of the class adds a significant amount of boiler plate and complexity and make the previous move hard to follow.
author Pierre-Yves David <pierre-yves.david@octobus.net>
date Mon, 20 Jan 2025 13:00:21 +0100
parents 70306aefa52b
children 7f848cfc4286
files mercurial/streamclone.py
diffstat 1 files changed, 24 insertions(+), 13 deletions(-) [+]
line wrap: on
line diff
--- a/mercurial/streamclone.py	Mon Jan 20 12:32:37 2025 +0100
+++ b/mercurial/streamclone.py	Mon Jan 20 13:00:21 2025 +0100
@@ -18,6 +18,7 @@
     Optional,
     Set,
     Tuple,
+    Type,
 )
 
 from .i18n import _
@@ -1134,21 +1135,30 @@
 ]
 
 
-def _v2_file_chunk(
-    fp: bundle2mod.unbundlepart,
-    data_len: int,
-    progress: scmutil.progress,
-    report: V2Report,
-) -> FileChunksT:
+class _FileChunker:
     """yield the chunk that constitute a file
 
-    This function exists as the counterpart of a future threaded version and
+    This class exists as the counterpart of a future threaded version and
     would not be very useful on its own.
     """
-    for chunk in util.filechunkiter(fp, limit=data_len):
-        report.byte_count += len(chunk)
-        progress.increment(step=len(chunk))
-        yield chunk
+
+    def __init__(
+        self,
+        fp: bundle2mod.unbundlepart,
+        data_len: int,
+        progress: scmutil.progress,
+        report: V2Report,
+    ):
+        self.fp = fp
+        self.report = report
+        self.progress = progress
+        self._chunks = util.filechunkiter(fp, limit=data_len)
+
+    def __iter__(self) -> FileChunksT:
+        for chunk in self._chunks:
+            self.report.byte_count += len(chunk)
+            self.progress.increment(step=len(chunk))
+            yield chunk
 
 
 def _v2_parse_files(
@@ -1157,6 +1167,7 @@
     file_count: int,
     progress: scmutil.progress,
     report: V2Report,
+    file_chunker: Type[_FileChunker] = _FileChunker,
 ) -> Iterator[FileInfoT]:
     """do the "stream-parsing" part of stream v2
 
@@ -1174,8 +1185,8 @@
             repo.ui.debug(
                 b'adding [%s] %s (%s)\n' % (src, name, util.bytecount(datalen))
             )
-
-        yield (src, name, _v2_file_chunk(fp, datalen, progress, report))
+        chunks = file_chunker(fp, datalen, progress, report)
+        yield (src, name, iter(chunks))
 
 
 def _write_files(vfsmap, info: Iterable[FileInfoT]):