import math
import weakref
from multiprocessing import shared_memory

import numpy as np


class SharedArray:
    def __init__(self, shape, dtype, name=None):
        if name:
            self.shm = shared_memory.SharedMemory(name=name)
        else:
            size = math.prod(shape) * np.dtype(dtype).itemsize
            self.shm = shared_memory.SharedMemory(create=True, size=size)
        self.arr = np.ndarray(shape, dtype, self.shm.buf)
        # This unlinks the shared memory buffer, but it will survive until the last
        # process closes up their file pointer. This could cause a problem if a
        # SharedArray is serialize and the Python object goes out of scope before
        # trying to deserialize it (in the same or a different process).
        weakref.finalize(self.arr, self.close)

    @property
    def name(self):
        return self.shm.name

    @property
    def array(self):
        return self.arr

    def result(self):
        return self.arr

    def close(self):
        self.arr = None
        try:
            self.shm.unlink()
        except FileNotFoundError:
            pass

    def __getstate__(self):
        return (self.arr.shape, self.arr.dtype.str, self.shm.name)

    def __setstate__(self, args):
        self.__init__(*args)
