# Path: taskloaf/shmem.py
# @contextlib.contextmanager
# def alloc_shmem(size, filepath):
#     assert not os.path.exists(filepath)
#     try:
#         init_shmem_file(filepath, size)
#         logger.info(f"creating shm {filepath}")
#         yield filepath
#     finally:
#         try:
#             logger.info(f"deleting shm {filepath}")
#             os.remove(filepath)
#         except FileNotFoundError:
#             logger.info(
#                 f"failed to delete shm {filepath} because it no longer exists"
#             )
#             pass
# 
# class Shmem:
#     def __init__(self, filepath, track_refs=False):
#         self.filepath = filepath
#         self.track_refs = track_refs
# 
#     def __enter__(self):
#         self.file = open(self.filepath, "r+b")
#         self.size = get_size_from_fd(self.file.fileno())
#         self.mmap = mmap_full_file(self.file.fileno())
# 
#         if not self.track_refs:
#             # mmap tracks references internally. This is annoying when trying
#             # to delete the mmap.  This code is some crooked trickery using
#             # ctypes and grabbing the pointer from a numpy buffer to create a
#             # memoryview from the mmap without mmap knowing about it so that
#             # the mmap can be closed without tracking its references (this WILL
#             # cause seg faults if there are existing references to the mmap
#             # segment when it's deleted)
#             # When using Shmem blocks through the allocator system, this is
#             # fine since then taskloaf performs its own memory tracking
#             temp_np = np.frombuffer(self.mmap, dtype=np.uint8)
#             ptr = temp_np.ctypes.data
#             del temp_np
#             ptrc = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_byte))
#             new_array = np.ctypeslib.as_array(ptrc, shape=(self.size,))
#             self.mem = memoryview(new_array.data.cast("B"))
#         else:
#             self.mem = memoryview(self.mmap)
# 
#         return self
# 
#     def __exit__(self, exc_type, exc_value, traceback):
#         self.file.close()
#         self.mmap.close()
# 
# Path: taskloaf/timer.py
# class Timer(object):
#     def __init__(self, output_fnc=print):
#         if output_fnc is None:
#             output_fnc = ignore_args
#         self.output_fnc = output_fnc
#         self.restart()
# 
#     def restart(self):
#         self.start = get_time()
# 
#     def report(self, name, should_restart=True):
#         text = name + " took " + str(get_time() - self.start)
#         self.output_fnc(text)
#         if should_restart:
#             self.restart()


import os
import time
import numpy as np
import multiprocessing
from taskloaf.shmem import alloc_shmem, Shmem
from taskloaf.timer import Timer


def sum_shmem(filepath):
    with Shmem(filepath) as sm:
        return np.sum(np.frombuffer(sm.mem))

def test_shmem():
    A = np.random.rand(100)

# Path: taskloaf/shmem.py
# @contextlib.contextmanager
# def alloc_shmem(size, filepath):
#     assert not os.path.exists(filepath)
#     try:
#         init_shmem_file(filepath, size)
#         logger.info(f"creating shm {filepath}")
#         yield filepath
#     finally:
#         try:
#             logger.info(f"deleting shm {filepath}")
#             os.remove(filepath)
#         except FileNotFoundError:
#             logger.info(
#                 f"failed to delete shm {filepath} because it no longer exists"
#             )
#             pass
# 
# class Shmem:
#     def __init__(self, filepath, track_refs=False):
#         self.filepath = filepath
#         self.track_refs = track_refs
# 
#     def __enter__(self):
#         self.file = open(self.filepath, "r+b")
#         self.size = get_size_from_fd(self.file.fileno())
#         self.mmap = mmap_full_file(self.file.fileno())
# 
#         if not self.track_refs:
#             # mmap tracks references internally. This is annoying when trying
#             # to delete the mmap.  This code is some crooked trickery using
#             # ctypes and grabbing the pointer from a numpy buffer to create a
#             # memoryview from the mmap without mmap knowing about it so that
#             # the mmap can be closed without tracking its references (this WILL
#             # cause seg faults if there are existing references to the mmap
#             # segment when it's deleted)
#             # When using Shmem blocks through the allocator system, this is
#             # fine since then taskloaf performs its own memory tracking
#             temp_np = np.frombuffer(self.mmap, dtype=np.uint8)
#             ptr = temp_np.ctypes.data
#             del temp_np
#             ptrc = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_byte))
#             new_array = np.ctypeslib.as_array(ptrc, shape=(self.size,))
#             self.mem = memoryview(new_array.data.cast("B"))
#         else:
#             self.mem = memoryview(self.mmap)
# 
#         return self
# 
#     def __exit__(self, exc_type, exc_value, traceback):
#         self.file.close()
#         self.mmap.close()
# 
# Path: taskloaf/timer.py
# class Timer(object):
#     def __init__(self, output_fnc=print):
#         if output_fnc is None:
#             output_fnc = ignore_args
#         self.output_fnc = output_fnc
#         self.restart()
# 
#     def restart(self):
#         self.start = get_time()
# 
#     def report(self, name, should_restart=True):
#         text = name + " took " + str(get_time() - self.start)
#         self.output_fnc(text)
#         if should_restart:
#             self.restart()


# Path: tests/test_shmem.py
import os
import time
import numpy as np
import multiprocessing
from taskloaf.shmem import alloc_shmem, Shmem
from taskloaf.timer import Timer



def sum_shmem(filepath):
    with Shmem(filepath) as sm:
        return np.sum(np.frombuffer(sm.mem))

def test_shmem():
    A = np.random.rand(100)

    with alloc_shmem(A.nbytes, 'test') as filepath: