1
\$\begingroup\$

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.

\$\endgroup\$
1
  • \$\begingroup\$ I updated the code with the possibility to write to files. I am thinking that maybe I can specify the opening mode: r, w, rw, or even append, and have some sort of "standard" interface. \$\endgroup\$ Jun 30, 2021 at 14:12

1 Answer 1

3
\$\begingroup\$
  • 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
\$\endgroup\$
4
  • 1
    \$\begingroup\$ Thank you, very clear and comprehensive! "Pickled objects are already length-aware." - today I learned! This is extremely neat, since now different pickled objects can simply be concatenated; the file pointer is advanced "for free". \$\endgroup\$ Jun 30, 2021 at 17:17
  • \$\begingroup\$ In general I discourage people from pickling any custom classes for long-term storage. While this might work with your particular code right now, you will possibly struggle to load files saved with your library a few versions ago. E.g. if you save them as object of type Metadata, you won't be able to load them if you rename the class to MetaData. Or even worse, pickle with mypackage.utils.Metadata and then do some refactoring - suddenly you'll find out that you're unable to unpickle it just because class has been moved to mypackage.types.Metadata! \$\endgroup\$ Mar 7, 2023 at 15:16
  • \$\begingroup\$ @BIOStheZerg responsible versioning and strong types are not mutually exclusive. Sit on that one for a minute. A responsible versioning system is not one that "allows anything" in serialization, but rather, has well-defined types for every version that has been published. \$\endgroup\$
    – Reinderien
    Mar 7, 2023 at 20:47
  • \$\begingroup\$ But that ("for every version") is the problem, not the solution. One should have data storage independent of the code, and a new version of the code should just have an updated readers/writers with as much backward (and forward!) compatibility as practical. But let's not get into a discussion in here :) \$\endgroup\$ Mar 11, 2023 at 10:32

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.