utils: proper file interface for HashedStreamReader (#5508)

This commit is contained in:
Batuhan Taskaya
2021-02-23 02:47:36 +02:00
committed by GitHub
parent 27c2401f59
commit d587d41d6c
2 changed files with 35 additions and 4 deletions
+18 -3
View File
@@ -1,6 +1,8 @@
import hashlib
import io
from funcy import cached_property
from dvc.hash_info import HashInfo
from dvc.istextfile import istextblock
from dvc.utils import dos2unix
@@ -61,17 +63,30 @@ class IterStream(io.RawIOBase):
return self.leftover[:n]
class HashedStreamReader:
class HashedStreamReader(io.IOBase):
PARAM_CHECKSUM = "md5"
def __init__(self, fobj):
self.fobj = fobj
self.md5 = hashlib.md5()
self.is_text_file = None
self.reader = fobj.read1 if hasattr(fobj, "read1") else fobj.read
super().__init__()
def readable(self):
return True
def tell(self):
return self.fobj.tell()
@cached_property
def _reader(self):
if hasattr(self.fobj, "read1"):
return self.fobj.read1
return self.fobj.read
def read(self, n=-1):
chunk = self.reader(n)
chunk = self._reader(n)
if self.is_text_file is None:
self.is_text_file = istextblock(chunk)
+17 -1
View File
@@ -9,7 +9,15 @@ def test_hashed_stream_reader(tmp_dir):
foo = tmp_dir / "foo"
with open(foo, "rb") as fobj:
stream_reader = HashedStreamReader(fobj)
assert stream_reader.read(3) == b"foo"
assert stream_reader.readable()
assert not stream_reader.seekable()
assert stream_reader.read(2) == b"fo"
assert stream_reader.tell() == 2
assert stream_reader.read(1) == b"o"
assert stream_reader.tell() == 3
hex_digest = file_md5(foo, LocalFileSystem(None, {}))
assert stream_reader.is_text_file
@@ -20,13 +28,21 @@ def test_hashed_stream_reader_as_chunks(tmp_dir):
tmp_dir.gen({"foo": b"foo \x00" * 16})
foo = tmp_dir / "foo"
actual_size = len(foo.read_bytes())
with open(foo, "rb") as fobj:
stream_reader = HashedStreamReader(fobj)
total_read = 0
while True:
chunk = stream_reader.read(16)
total_read += len(chunk)
assert stream_reader.tell() == total_read
if not chunk:
break
assert stream_reader.tell() == actual_size == total_read
hex_digest = file_md5(foo, LocalFileSystem(None, {}))
assert not stream_reader.is_text_file
assert hex_digest == stream_reader.hash_info.value