mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-01 23:08:43 +01:00
97 lines
2.3 KiB
Python
97 lines
2.3 KiB
Python
"""Common utilities for serialize."""
|
|
import os
|
|
from contextlib import contextmanager
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Callable,
|
|
ContextManager,
|
|
Dict,
|
|
Protocol,
|
|
Union,
|
|
)
|
|
|
|
from funcy import reraise
|
|
|
|
from dvc.exceptions import DvcException
|
|
|
|
if TYPE_CHECKING:
|
|
from dvc.fs import FileSystem
|
|
from dvc.types import AnyPath
|
|
|
|
|
|
class DumperFn(Protocol):
|
|
def __call__(
|
|
self, path: "AnyPath", data: Any, fs: "FileSystem" = None
|
|
) -> Any:
|
|
...
|
|
|
|
|
|
class ModifierFn(Protocol):
|
|
def __call__(
|
|
self, path: "AnyPath", fs: "FileSystem" = None
|
|
) -> ContextManager[Dict]:
|
|
...
|
|
|
|
|
|
class LoaderFn(Protocol):
|
|
def __call__(self, path: "AnyPath", fs: "FileSystem" = None) -> Any:
|
|
...
|
|
|
|
|
|
ReadType = Union[bytes, None, str]
|
|
ParserFn = Callable[[ReadType, "AnyPath"], dict]
|
|
|
|
|
|
class ParseError(DvcException):
|
|
"""Errors while parsing files"""
|
|
|
|
def __init__(self, path: "AnyPath", message: str):
|
|
from dvc.utils import relpath
|
|
|
|
path = relpath(path)
|
|
self.path = path
|
|
super().__init__(f"unable to read: '{path}', {message}")
|
|
|
|
|
|
class EncodingError(ParseError):
|
|
"""We could not read a file with the given encoding"""
|
|
|
|
def __init__(self, path: "AnyPath", encoding: str):
|
|
self.encoding = encoding
|
|
super().__init__(path, f"is not valid {encoding}")
|
|
|
|
|
|
def _load_data(path: "AnyPath", parser: ParserFn, fs: "FileSystem" = None):
|
|
open_fn = fs.open if fs else open
|
|
encoding = "utf-8"
|
|
with open_fn(path, encoding=encoding) as fd: # type: ignore
|
|
with reraise(UnicodeDecodeError, EncodingError(path, encoding)):
|
|
return parser(fd.read(), path)
|
|
|
|
|
|
def _dump_data(
|
|
path,
|
|
data: Any,
|
|
dumper: DumperFn,
|
|
fs: "FileSystem" = None,
|
|
**dumper_args,
|
|
):
|
|
open_fn = fs.open if fs else open
|
|
with open_fn(path, "w+", encoding="utf-8") as fd: # type: ignore
|
|
dumper(data, fd, **dumper_args)
|
|
|
|
|
|
@contextmanager
|
|
def _modify_data(
|
|
path: "AnyPath",
|
|
parser: ParserFn,
|
|
dumper: DumperFn,
|
|
fs: "FileSystem" = None,
|
|
):
|
|
exists_fn = fs.exists if fs else os.path.exists
|
|
file_exists = exists_fn(path) # type: ignore
|
|
data = _load_data(path, parser=parser, fs=fs) if file_exists else {}
|
|
yield data
|
|
dumper(path, data, fs=fs)
|