dvc/utils/stream: make IterStream seekable and peekable (#5084)

fixes #5021

dvc/utils/stream.py

@@ -6,11 +6,14 @@ class IterStream(io.RawIOBase):
 
     def __init__(self, iterator):  # pylint: disable=super-init-not-called
         self.iterator = iterator
-        self.leftover = None
+        self.leftover = b""
 
     def readable(self):
         return True
 
+    def writable(self) -> bool:
+        return False
+
     # Python 3 requires only .readinto() method, it still uses other ones
     # under some circumstances and falls back if those are absent. Since
     # iterator already constructs byte strings for us, .readinto() is not the
@@ -38,8 +41,16 @@ class IterStream(io.RawIOBase):
 
         # Return an arbitrary number or bytes
         if n <= 0:
-            self.leftover = None
+            self.leftover = b""
             return chunk
 
         output, self.leftover = chunk[:n], chunk[n:]
         return output
+
+    def peek(self, n):
+        while len(self.leftover) < n:
+            try:
+                self.leftover += next(self.iterator)
+            except StopIteration:
+                break
+        return self.leftover[:n]

tests/unit/utils/test_http.py

@@ -5,6 +5,16 @@ import requests
 from dvc.utils.http import open_url
 
 
+def read_nbytes(fd, num, mode="r"):
+    data = b"" if mode == "rb" else ""
+    while len(data) < num:
+        chunk = fd.read(num - len(data))
+        if not chunk:
+            break
+        data += chunk
+    return data
+
+
 def test_open_url(tmp_path, monkeypatch, http):
     # Simulate bad connection
     original_iter_content = requests.Response.iter_content
@@ -27,5 +37,20 @@ def test_open_url(tmp_path, monkeypatch, http):
 
     with open_url((http / "sample.txt").url) as fd:
         # Test various .read() variants
-        assert fd.read(len(text)) == text
+        assert read_nbytes(fd, len(text), mode="r") == text
         assert fd.read() == text
+        assert fd.read() == ""
+
+
+def test_open_url_peek_rb(tmp_path, monkeypatch, http):
+    # Goes over seek feature in 'rb' mode
+    text = "0123456789" * (io.DEFAULT_BUFFER_SIZE // 10 + 1)
+    http.gen("sample.txt", text * 2)
+
+    with open_url((http / "sample.txt").url, mode="rb") as fd:
+        text = text.encode("utf8")
+        assert fd.peek(len(text)) == text
+        assert read_nbytes(fd, len(text), mode="rb") == text
+        assert fd.peek(len(text)) == text
+        assert fd.read() == text
+        assert fd.read() == b""