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)
