numpy_array.py

import logging

import numpy as np
from numpy.lib.format import open_memmap

from .registry import register_array_parser


@register_array_parser("numpy")
class NumpyArrayParser(object):
    def __init__(self):
        pass

    def read(self, path):
        logging.info("Reading from %s using numpy format" % path)
        arr = np.load(path, mmap_mode="r")
        logging.info("Done reading from %s" % path)
        return arr

    def write(self, path, arr):
        logging.info("Writing to %s using numpy format" % path)
        # np.save would load the entire memmap array up into CPU.  So we manually open
        # an empty npy file with memmap mode and manually flush it instead.
        new_arr = open_memmap(path, mode="w+", dtype=arr.dtype, shape=arr.shape)
        new_arr[:] = arr[:]
        logging.info("Done writing to %s" % path)