Class PartialFileIO

modelscope/hub/commit_scheduler.py:313–381

A file-like object that reads only the first part of a file.

class PartialFileIO(BytesIO):
    """A file-like object that reads only the first part of a file."""

    def __init__(self, file_path: Union[str, Path], size_limit: int) -> None:
        self._file_path = Path(file_path)
        self._file = None
        self._size_limit = size_limit
        self.open()

    def open(self) -> None:
        """Open the file and initialize size limit."""
        if self._file is not None:
            return
        try:
            self._file = self._file_path.open('rb')
            self._size_limit = min(
                self._size_limit or float('inf'),
                os.fstat(self._file.fileno()).st_size)
        except OSError as e:
            logger.error(f'Failed to open file {self._file_path}: {e}')
            raise

    def close(self) -> None:
        """Close the file if it's open."""
        if self._file is not None:
            self._file.close()
            self._file = None

    def __del__(self) -> None:
        self.close()
        return super().__del__()

    def __repr__(self) -> str:
        return f'<PartialFileIO file_path={self._file_path} size_limit={self._size_limit}>'

    def __len__(self) -> int:
        return self._size_limit

    def __getattribute__(self, name: str):
        if name.startswith('_') or name in {
                'read', 'tell', 'seek', 'close', 'open'
        }:  # only 5 public methods supported
            return super().__getattribute__(name)
        raise NotImplementedError(f"PartialFileIO does not support '{name}'.")

    def tell(self) -> int:
        return self._file.tell()

    def seek(self, __offset: int, __whence: int = SEEK_SET) -> int:
        """Seek to a position in the file, but never beyond size_limit."""
        if __whence == SEEK_END:
            __offset = len(self) + __offset
            __whence = SEEK_SET

        pos = self._file.seek(__offset, __whence)
        if pos > self._size_limit:
            return self._file.seek(self._size_limit)
        return pos