afnio.serialization

def save(
    obj: object,
    f: FILE_LIKE,
    pickle_protocol: int = DEFAULT_PROTOCOL,
) -> None:
    """
    Saves an object to a disk file using zip compression and pickle serialization.

    Args:
        obj: The object to be saved.
        f (Union[str, os.PathLike, BinaryIO, IO[bytes]]): A file-like object (must
            implement write/flush) or a string or `os.PathLike` object containing
            a file name.
        pickle_protocol: Pickle protocol version.

    Note:
        A common Afnio convention is to save variables using `.hf` file extension.

        The `.hf` extension is a naming convention inspired by the chemical symbol
        for [Hafnium (Hf)](https://en.wikipedia.org/wiki/Hafnium). "Afnio" is the
        Italian word for Hafnium.

    Examples:
        >>> # Save to file
        >>> x = afnio.Variable(data="You are a doctor.", role="system prompt")
        >>> afnio.save(x, 'variable.hf')
        >>> # Save to io.BytesIO buffer
        >>> buffer = io.BytesIO()
        >>> afnio.save(x, buffer)
    """
    _check_save_filelike(f)

    with _open_zipfile_writer(f) as opened_zipfile:
        _save(
            obj,
            opened_zipfile,
            pickle_protocol,
        )
        return

def load(f: FILE_LIKE) -> Any:
    """
    Loads an object from a disk file using zip compression and pickle serialization.

    Args:
        f (Union[str, os.PathLike, BinaryIO, IO[bytes]]): A file-like object (must
            implement `read`) or a string or `os.PathLike` object containing
            a file name.

    Returns:
        The deserialized object.

    Examples:
        >>> # Load from file
        >>> obj = afnio.load('model.hf')
        >>> # Load from io.BytesIO buffer
        >>> buffer = io.BytesIO()
        >>> obj = afnio.load(buffer)
    """
    with _open_zipfile_reader(f) as zip_reader:
        if "data.pkl" not in zip_reader.namelist():
            raise RuntimeError(
                "Missing 'data.pkl' in archive. File might be corrupted."
            )

        # Read the serialized object
        with zip_reader.open("data.pkl", "r") as f:
            obj = pickle.load(f)

    return obj