mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
utils: proper file interface for HashedStreamReader (#5508)
This commit is contained in:
+18
-3
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user