Python class for handling custom file type

Question

Background

I am devising a simple file format to store some (possibly large) binary data together with some metadata that specifies how the data was generated:

<HEADER | 4 bytes> (specifies 'size' of METADATA)
<METADATA | 'size' bytes>
<DATA>

The use case is that if at some later point I just want to inspect the metadata (i.e. without touching the actual data), then I won't read the entire structure, but only the <METADATA> part (of known size).

I implemented the following class to facilitate working with such files. It behaves similarly to a singleton, loads metadata or data when first requested, then holds the reference to them during the entire lifespan of the class.

I am planning to use hypothesis for testing (generating random data of various types).

Code

import numpy as np
import pickle
from pathlib import Path
from typing import Any, Union


class MyFile:

    HEADER_SIZE_BYTES = 4

    def __init__(self, filename: Union[str, Path]):
        self.filename = filename

        self.__metadata = None
        self.__data = None
        
        with open(filename, 'rb') as fp:
            self.metadata_size = int.from_bytes(fp.read(self.HEADER_SIZE_BYTES), "little")

    def get_metadata(self) -> Any:
        if self.__metadata is None:
            with open(self.filename, "rb") as fp:
                fp.seek(self.HEADER_SIZE_BYTES, 1)
                self.__metadata = pickle.loads(fp.read(self.metadata_size))

        return self.__metadata

    def set_metadata(self, metadata: Any) -> None:
        self.__metadata = metadata

    def get_data(self) -> Any:
        if self.__data is None:
            with open(self.filename, "rb") as fp:
                fp.seek(self.HEADER_SIZE_BYTES + self.metadata_size, 1)
                self.__data = pickle.loads(fp.read())

        return self.__data
    
    def set_data(self, data: Any) -> None:
        self.__data = data

    def write_all(self) -> None:
        if self.__metadata is None:
            raise ValueError("No metadata to write!")

        if self.__data is None:
            raise ValueError("No data to write!")

        metadata_bytes = pickle.dumps(self.__metadata)
        data_bytes = pickle.dumps(self.__data)
        header_bytes = len(metadata_bytes).to_bytes(self.HEADER_SIZE_BYTES, "little")

        with open(self.filename, 'wb') as fp:
            fp.write(header_bytes)
            fp.write(metadata_bytes)
            fp.write(data_bytes)


if __name__ == "__main__":
    f = MyFile("out.epkl")
    
    meta_w = {
        "foo": "bar",
        "baz": [1, 2, 3],
    }
    data_w = np.random.rand(3, 10, 10)

    f.set_metadata(meta_w)
    f.set_data(data_w)
    f.write_all()

    g = MyFile("out.epkl")
    
    meta_r = g.get_metadata()
    data_r = g.get_data()
    
    assert meta_r == meta_w
    assert np.allclose(data_r, data_w)

Comments

What are the crucial aspects that you, as a user of this library, would require from it?

I feel like things can be improved: there's some duplicated code, and possibly some edge-cases that I did not consider.

Answer 1

• It's not necessary to represent this as a class; a read and write function will suffice
• Don't use double-underscores; that's reserved for name mangling
• The second argument of Union[str, Path] is more generically expressed as os.PathLike
• Pickled objects are already length-aware. The much easier approach than what you've shown here is to omit any manual length information, and simply do two sequential serialization calls to pickle.
• Particularly since you're using pickle and not e.g. JSON, there is no excuse to have poorly-typed data. Replace your dictionary with dataclass or similar.

The following suggested code represents the write operation as a naive method assuming both metadata and body are already in memory; and read as an iterator that yields them one at a time. Drop the logging if you don't need it; it's more to demonstrate successful operation than anything else.

import logging
import pickle
from dataclasses import dataclass
from os import PathLike
from sys import stdout
from typing import Iterable, Union, List


import numpy as np

logger = logging.getLogger('pickling_experments')


@dataclass
class Metadata:
    foo: str
    baz: List[int]


def write_pickle(
    metadata: Metadata, 
    body: np.ndarray,
    path: Union[str, PathLike] = 'out.epkl',
) -> None:
    with open(path, 'wb') as f:
        pickle.dump(metadata, f)
        logger.debug(f'Position after metadata: {f.tell()}')
        pickle.dump(body, f)
        logger.debug(f'Position after {body.nbytes}-byte body: {f.tell()}')


def read_pickle(
    path: Union[str, PathLike] = 'out.epkl',
    get_metadata: bool = True,
    get_body: bool = True,
) -> Iterable[
    Union[Metadata, np.ndarray]
]:
    with open(path, 'rb') as f:
        metadata = pickle.load(f)
        logger.debug(f'Position after metadata: {f.tell()}')
        if get_metadata:
            yield metadata

        if not get_body:
            return
        body = pickle.load(f)
        logger.debug(f'Position after {body.nbytes}-byte body: {f.tell()}')

    yield body


def test():
    logger.addHandler(logging.StreamHandler(stdout))
    logger.setLevel(logging.DEBUG)

    print('Write:')
    write_pickle(
        Metadata(foo='bar', baz=[1, 2, 3]), 
        np.random.random((512, 512)),
    )

    print('\nCombined read:', flush=True)
    metadata, body = read_pickle()

    print('\nSeparated read:', flush=True)
    results = read_pickle()
    metadata = next(results)
    print('Twiddle your thumbs for a while here; large memory not yet consumed', flush=True)
    body = next(results)

    print('\nMetadata-only read:', flush=True)
    metadata, = read_pickle(get_body=False)


if __name__ == '__main__':
    test()

Output:

Write:
Position after metadata: 68
Position after 2097152-byte body: 2097383

Combined read:
Position after metadata: 68
Position after 2097152-byte body: 2097383

Separated read:
Position after metadata: 68
Twiddle your thumbs for a while here; large memory not yet consumed
Position after 2097152-byte body: 2097383

Metadata-only read:
Position after metadata: 68