diff --git a/dvc/data/db/reference.py b/dvc/data/db/reference.py index 01589f49c..36877d6ed 100644 --- a/dvc/data/db/reference.py +++ b/dvc/data/db/reference.py @@ -1,4 +1,3 @@ -import io import logging import os from typing import TYPE_CHECKING, Dict @@ -37,10 +36,10 @@ class ReferenceObjectDB(ObjectDB): pass fs_path = self.hash_to_path(hash_info.value) try: - with self.fs.open(fs_path, "rb") as fobj: - ref_file = ReferenceHashFile.from_bytes( - fobj.read(), fs_cache=self._fs_cache - ) + ref_file = ReferenceHashFile.from_bytes( + self.fs.cat_file(fs_path), + fs_cache=self._fs_cache, + ) except OSError: raise FileNotFoundError try: @@ -70,10 +69,8 @@ class ReferenceObjectDB(ObjectDB): ) ref_file = ReferenceHashFile(from_info, from_fs, hash_info) self._obj_cache[hash_info] = ref_file - ref_fobj = io.BytesIO(ref_file.to_bytes()) - ref_fobj.seek(0) try: - self.fs.upload(ref_fobj, to_info) + self.fs.pipe_file(to_info, ref_file.to_bytes()) except OSError as exc: if isinstance(exc, FileExistsError) or ( os.name == "nt" diff --git a/dvc/data/tree.py b/dvc/data/tree.py index 0d60327ca..1cbdb9e2e 100644 --- a/dvc/data/tree.py +++ b/dvc/data/tree.py @@ -62,8 +62,7 @@ class Tree(HashFile): memfs = MemoryFileSystem() fs_path = "memory://{}".format(tmp_fname("")) - with memfs.open(fs_path, "wb") as fobj: - fobj.write(self.as_bytes()) + memfs.pipe_file(fs_path, self.as_bytes()) self.fs = memfs self.fs_path = fs_path if hash_info: diff --git a/dvc/fs/base.py b/dvc/fs/base.py index 0ca3d32f0..cd7102acb 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -188,6 +188,54 @@ class FileSystem: ) -> "IO": # pylint: disable=arguments-differ return self.fs.open(path, mode=mode, encoding=encoding, **kwargs) + def cat( + self, + path: Union[AnyFSPath, List[AnyFSPath]], + recursive: bool = False, + on_error: "Literal['raise', 'omit', 'return']" = "raise", + **kwargs: Any, + ) -> Union[bytes, Dict[AnyFSPath, bytes]]: + return self.fs.cat( + path, recursive=recursive, on_error=on_error, **kwargs + ) + + def cat_ranges( + self, + paths: List[AnyFSPath], + starts: List[int], + ends: List[int], + max_gap: int = None, + **kwargs, + ) -> List[bytes]: + return self.fs.cat_ranges( + paths, starts, ends, max_gap=max_gap, **kwargs + ) + + def cat_file( + self, + path: AnyFSPath, + start: int = None, + end: int = None, + **kwargs: Any, + ) -> bytes: + return self.fs.cat_file(path, start=start, end=end, **kwargs) + + def pipe_file(self, path: AnyFSPath, value: bytes, **kwargs: Any) -> None: + return self.fs.pipe_file(path, value, **kwargs) + + def pipe( + self, + path: Union[AnyFSPath, Dict[AnyFSPath, bytes]], + value: Optional[bytes] = None, + **kwargs: Any, + ) -> None: + return self.fs.pipe(path, value=value, **kwargs) + + def touch( + self, path: AnyFSPath, truncate: bool = True, **kwargs: Any + ) -> None: + return self.fs.touch(path, truncate=truncate, **kwargs) + def checksum(self, path: AnyFSPath) -> str: return self.fs.checksum(path) diff --git a/dvc/fs/memory.py b/dvc/fs/memory.py index 436cdac36..1b6130720 100644 --- a/dvc/fs/memory.py +++ b/dvc/fs/memory.py @@ -10,8 +10,6 @@ from .base import FileSystem class MemoryFileSystem(FileSystem): # pylint:disable=abstract-method scheme = Schemes.MEMORY PARAM_CHECKSUM = "md5" - TRAVERSE_PREFIX_LEN = 2 - DEFAULT_BLOCKSIZE = 4096 def __eq__(self, other): # NOTE: all fsspec MemoryFileSystem instances are equivalent and use a @@ -26,8 +24,3 @@ class MemoryFileSystem(FileSystem): # pylint:disable=abstract-method from fsspec.implementations.memory import MemoryFileSystem as MemFS return MemFS(**self.fs_args) - - def open(self, *args, **kwargs): - with super().open(*args, **kwargs) as fobj: - fobj.blocksize = self.DEFAULT_BLOCKSIZE - return fobj diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index 27a12e415..fe2227fad 100644 --- a/dvc/repo/experiments/init.py +++ b/dvc/repo/experiments/init.py @@ -177,8 +177,7 @@ def init_deps(stage: PipelineStage) -> List["Dependency"]: localfs.makedirs(path) for path in files: localfs.makedirs(localfs.path.parent(path), exist_ok=True) - with localfs.open(path, "w", encoding="utf-8"): - pass + localfs.touch(path) return new_deps