LazyCombiner

class LazyCombiner(LazyBuffer):
    def __init__(
        self,
        buffer_or_path: str | UPath | BufferWriter,
        *,
        mode: Literal["a", "w"] = "w",
        packer: type[LazyCodec] | LazyCodec | None = None,
        fs: FileSystem | None = None,
    ):
        """
        The mode resembles typical mode designations and implies the same meaning.
        If the mode is 'w', the file is overwritten.
        If the mode is 'a', the file is appended.

        :param buffer_or_path: target buffer or file path
        :param mode: mode of operation, 'w' for write and 'a' for append
        :param packer: packer object to be used for packing the object
        :param fs: `FileSystem` object to be used for storing
        """
        super().__init__(buffer_or_path, packer, fs)

        self._mode: str = mode

        self._toc: dict | list = None  # type: ignore

    def __enter__(self):
        if isinstance(self._buffer_or_path, str):
            # we do not use a unified interface via `LocalFileSystem`
            # because for local storage we can skip the cache file
            if self._fs:
                self._buffer = TemporaryFile()
                if self._fs.exists(self._buffer_or_path):
                    with self._fs.open(self._buffer_or_path, "rb") as s3_file:
                        while chunk := s3_file.read(config.read_buffer_size):
                            self._buffer.write(chunk)
                    self._buffer.seek(0)
            else:
                mode: str = (
                    "wb"
                    if not os.path.exists(self._buffer_or_path) or self._mode == "w"
                    else "r+b"
                )
                self._buffer = open(  # type: ignore
                    self._buffer_or_path, mode, config.write_buffer_size
                )
        elif isinstance(self._buffer_or_path, UPath):
            self._buffer = self._buffer_or_path.open(
                "wb"
                if not self._buffer_or_path.exists() or self._mode == "w"
                else "r+b",
                config.write_buffer_size,
            )
            if not self._buffer.seekable():
                raise ValueError(
                    f"The underlying filesystem of the given UPath ({self._buffer_or_path}) does not support random write."
                )
                self._unseekable_upath = True
                self._buffer.close()
                self._buffer = TemporaryFile()
                if self._buffer_or_path.exists():
                    with self._buffer_or_path.open(
                        "rb", config.read_buffer_size
                    ) as source:
                        _copy_obj(source, self._buffer)
        elif isinstance(self._buffer_or_path, (BytesIO, BufferedReader)):
            self._buffer = self._buffer_or_path
            if self._mode == "a":
                # need to read the header anyway
                self._buffer.seek(0)
        else:
            raise ValueError("Expecting a buffer or path.")

        if self._mode == "w":
            self._buffer.write(LazyWriter.magic)
            self._header_start = self._buffer.tell()
            self._buffer.write(b"\0" * 20)
            self._file_start = self._buffer.tell()
        else:
            sep_a, sep_b, sep_c = (
                LazyWriter.magic_len(),
                LazyWriter.magic_len() + 10,
                LazyWriter.magic_len() + 20,
            )

            ini_position: int = self._buffer.tell()
            header: bytes = self._buffer.read(sep_c)

            def _raise_invalid(msg: str):
                self._buffer.seek(ini_position)
                raise ValueError(msg)

            if not config.check_compatibility(header[:sep_a]):
                _raise_invalid(
                    "Invalid file format, cannot append to the current file."
                )

            toc_start: int = self._packer.decode(header[sep_a:sep_b].lstrip(b"\0"))
            toc_size: int = self._packer.decode(header[sep_b:sep_c].lstrip(b"\0"))

            self._buffer.seek(ini_position + sep_c + toc_start)
            self._toc = self._packer.decode(self._buffer.read(toc_size)).get("t", None)

            if isinstance(self._toc, list):
                if any(not isinstance(i, int) for i in self._toc):
                    _raise_invalid("The given file is not a valid combined file.")
            elif isinstance(self._toc, dict):
                if any(not isinstance(i, int) for i in self._toc.values()):
                    _raise_invalid("The given file is not a valid combined file.")
            else:
                _raise_invalid("The given file is not a valid combined file.")

            self._header_start = ini_position + sep_a
            self._file_start = ini_position + sep_c
            self._buffer.seek(ini_position + sep_c + toc_start)

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._finalize({"t": self._toc})

        self.cleanup()

    def write(self, obj: Generator, name: str | None = None) -> None:
        """
        Write the object to the file.

        :param obj: the object to be written to the file, the generator shall yield data
        :param name: a name to be assigned to the object, only required when combining in dict mode
        """
        if self._toc is None:
            self._toc = [] if name is None else {}

        if name is None:
            if not isinstance(self._toc, list):
                raise ValueError("Need a name when combining in dict mode.")
        else:
            if not isinstance(self._toc, dict):
                raise ValueError("Cannot assign a name when combining in list mode.")
            if name in self._toc:
                raise ValueError(f"File {name} already exists.")

        start: int = self._buffer.tell() - self._file_start
        for chunk in obj:
            self._buffer.write(chunk)

        if name is None:
            self._toc.append(start)
        else:
            self._toc[name] = start