diff --git a/dvc/utils/stream.py b/dvc/utils/stream.py index cfa18d183..c2d3fe33e 100644 --- a/dvc/utils/stream.py +++ b/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] diff --git a/tests/test.p12 b/tests/test.p12 new file mode 100644 index 000000000..6c391b054 Binary files /dev/null and b/tests/test.p12 differ diff --git a/tests/unit/remote/test.p12 b/tests/unit/remote/test.p12 new file mode 100644 index 000000000..6c391b054 Binary files /dev/null and b/tests/unit/remote/test.p12 differ diff --git a/tests/unit/utils/test_http.py b/tests/unit/utils/test_http.py index a2f97701f..2051bcbaf 100644 --- a/tests/unit/utils/test_http.py +++ b/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""