# Path: caput/mpiutil.py
# def enable_mpi_exception_handler():
#     def mpi_excepthook(exc_type, exc_obj, exc_tb):
#     def __repr__(self):
# def active_comm(aprocs):
# def active(aprocs):
# def close(aprocs):
# def partition_list(full_list, i, n, method="con"):
#     def _partition(N, n, i):
# def partition_list_mpi(full_list, method="con", comm=_comm):
# def mpirange(*args, **kargs):
# def barrier(comm=_comm):
# def bcast(data, root=0, comm=_comm):
# def allreduce(sendobj, op=None, comm=_comm):
# def parallel_map(func, glist, root=None, method="con", comm=_comm):
# def typemap(dtype):
# def split_m(n, m):
# def split_all(n, comm=_comm):
# def split_local(n, comm=_comm):
# def gather_local(global_array, local_array, local_start, root=0, comm=_comm):
# def transpose_blocks(row_array, shape, comm=_comm):
# def allocate_hdf5_dataset(fname, dsetname, shape, dtype, comm=_comm):
# def lock_and_write_buffer(obj, fname, offset, size):
# def parallel_rows_write_hdf5(fname, dsetname, local_data, shape, comm=_comm):
#     def __init__(
#         self,
#         add_mpi_info: bool = True,
#         level_rank0: int = logging.INFO,
#         level_all: int = logging.WARN,
#         comm: Optional["MPI.Intracomm"] = None,
#     ):
#     def filter(self, record):
#     def __init__(self, self_module, baked_args=None):
#     def __getattr__(self, name):
#     def __call__(self, **kwargs):
# class _close_message:
# class MPILogFilter(logging.Filter):  # pylint: disable=too-few-public-methods
# class SelfWrapper(ModuleType):
#     N = len(full_list)
# 
# Path: caput/mpiarray.py
# class _global_resolver:
# class MPIArray(np.ndarray):
# class DummyContext:
#     def __init__(self, array):
#     def _resolve_slice(self, slobj):
#     def __getitem__(self, slobj):
#     def __setitem__(self, slobj, value):
#     def global_shape(self):
#     def global_shape(self, var):
#     def axis(self):
#     def local_shape(self):
#     def local_shape(self, var):
#     def local_offset(self):
#     def local_offset(self, var):
#     def local_array(self):
#     def comm(self):
#     def __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs):
#     def global_slice(self):
#     def wrap(cls, array, axis, comm=None):
#     def redistribute(self, axis):
#     def enumerate(self, axis):
#     def from_hdf5(cls, f, dataset, comm=None, axis=0, sel=None):
#     def to_hdf5(
#         self,
#         f,
#         dataset,
#         create=False,
#         chunks=None,
#         compression=None,
#         compression_opts=None,
#     ):
#     def transpose(self, *axes):
#     def reshape(self, *shape):
#     def copy(self):
#     def gather(self, rank=0):
#     def allgather(self):
#     def _to_hdf5_serial(self, filename, dataset, create=False):
#     def _partition_io(self, skip=False, threshold=1.99):
# def _partition_sel(sel, split_axis, n, slice_):
# def _len_slice(slice_, n):
# def _reslice(slice_, n, subslice):
# def _expand_sel(sel, naxis):
#     def __enter__(self):
#     def __exit__(self, *exc):


import os
import unittest
import numpy as np
        import h5py
from caput import mpiutil, mpiarray

            with h5py.File(fname, "r") as f:

                h5ds = f["testds"][:]

                assert (h5ds == ga).all()

        ds2 = mpiarray.MPIArray.from_hdf5(fname, "testds")

        assert (ds2 == ds).all()

        mpiutil.barrier()

        # Check that reading over another distributed axis works
        ds3 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1)
        assert ds3.shape[0] == gshape[0]
        assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0]
        ds3 = ds3.redistribute(axis=0)
        assert (ds3 == ds).all()
        mpiutil.barrier()

        # Check a read with an arbitrary slice in there. This only checks the shape is correct.
        ds4 = mpiarray.MPIArray.from_hdf5(
            fname, "testds", axis=1, sel=(np.s_[3:10:2], np.s_[1:16:3])
        )
        assert ds4.shape[0] == 4
        assert ds4.shape[1] == mpiutil.split_local(5)[0]
        mpiutil.barrier()

        # Check the read with a slice along the axis being read

# Path: caput/mpiutil.py
# def enable_mpi_exception_handler():
#     def mpi_excepthook(exc_type, exc_obj, exc_tb):
#     def __repr__(self):
# def active_comm(aprocs):
# def active(aprocs):
# def close(aprocs):
# def partition_list(full_list, i, n, method="con"):
#     def _partition(N, n, i):
# def partition_list_mpi(full_list, method="con", comm=_comm):
# def mpirange(*args, **kargs):
# def barrier(comm=_comm):
# def bcast(data, root=0, comm=_comm):
# def allreduce(sendobj, op=None, comm=_comm):
# def parallel_map(func, glist, root=None, method="con", comm=_comm):
# def typemap(dtype):
# def split_m(n, m):
# def split_all(n, comm=_comm):
# def split_local(n, comm=_comm):
# def gather_local(global_array, local_array, local_start, root=0, comm=_comm):
# def transpose_blocks(row_array, shape, comm=_comm):
# def allocate_hdf5_dataset(fname, dsetname, shape, dtype, comm=_comm):
# def lock_and_write_buffer(obj, fname, offset, size):
# def parallel_rows_write_hdf5(fname, dsetname, local_data, shape, comm=_comm):
#     def __init__(
#         self,
#         add_mpi_info: bool = True,
#         level_rank0: int = logging.INFO,
#         level_all: int = logging.WARN,
#         comm: Optional["MPI.Intracomm"] = None,
#     ):
#     def filter(self, record):
#     def __init__(self, self_module, baked_args=None):
#     def __getattr__(self, name):
#     def __call__(self, **kwargs):
# class _close_message:
# class MPILogFilter(logging.Filter):  # pylint: disable=too-few-public-methods
# class SelfWrapper(ModuleType):
#     N = len(full_list)
# 
# Path: caput/mpiarray.py
# class _global_resolver:
# class MPIArray(np.ndarray):
# class DummyContext:
#     def __init__(self, array):
#     def _resolve_slice(self, slobj):
#     def __getitem__(self, slobj):
#     def __setitem__(self, slobj, value):
#     def global_shape(self):
#     def global_shape(self, var):
#     def axis(self):
#     def local_shape(self):
#     def local_shape(self, var):
#     def local_offset(self):
#     def local_offset(self, var):
#     def local_array(self):
#     def comm(self):
#     def __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs):
#     def global_slice(self):
#     def wrap(cls, array, axis, comm=None):
#     def redistribute(self, axis):
#     def enumerate(self, axis):
#     def from_hdf5(cls, f, dataset, comm=None, axis=0, sel=None):
#     def to_hdf5(
#         self,
#         f,
#         dataset,
#         create=False,
#         chunks=None,
#         compression=None,
#         compression_opts=None,
#     ):
#     def transpose(self, *axes):
#     def reshape(self, *shape):
#     def copy(self):
#     def gather(self, rank=0):
#     def allgather(self):
#     def _to_hdf5_serial(self, filename, dataset, create=False):
#     def _partition_io(self, skip=False, threshold=1.99):
# def _partition_sel(sel, split_axis, n, slice_):
# def _len_slice(slice_, n):
# def _reslice(slice_, n, subslice):
# def _expand_sel(sel, naxis):
#     def __enter__(self):
#     def __exit__(self, *exc):


# Path: caput/tests/test_mpiarray.py
import os
import unittest
import numpy as np
        import h5py
from caput import mpiutil, mpiarray


            with h5py.File(fname, "r") as f:

                h5ds = f["testds"][:]

                assert (h5ds == ga).all()

        ds2 = mpiarray.MPIArray.from_hdf5(fname, "testds")

        assert (ds2 == ds).all()

        mpiutil.barrier()

        # Check that reading over another distributed axis works
        ds3 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1)
        assert ds3.shape[0] == gshape[0]
        assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0]
        ds3 = ds3.redistribute(axis=0)
        assert (ds3 == ds).all()
        mpiutil.barrier()

        # Check a read with an arbitrary slice in there. This only checks the shape is correct.
        ds4 = mpiarray.MPIArray.from_hdf5(
            fname, "testds", axis=1, sel=(np.s_[3:10:2], np.s_[1:16:3])
        )
        assert ds4.shape[0] == 4
        assert ds4.shape[1] == mpiutil.split_local(5)[0]
        mpiutil.barrier()

        # Check the read with a slice along the axis being read

        ds5 = mpiarray.MPIArray.from_hdf5(