# Path: taskloaf/mpi.py
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()


import sys
import logging
import _pytest.runner
from mpi4py import MPI
from taskloaf.mpi import MPIComm


def pytest_runtest_protocol(item, nextitem):
    n_procs = getattr(item.function, 'n_procs', 1)
    rank = MPI.COMM_WORLD.Get_rank()
    size = MPI.COMM_WORLD.Get_size()


    if n_procs > size:
        print(
            'skipping ' + str(item.name) +
            ' because it needs ' + str(n_procs) + ' MPI procs'
        )
        return True

    if rank >= n_procs:
        new_comm = MPI.COMM_WORLD.Split(color = 0)
        MPI.COMM_WORLD.Barrier()
        return True
    else:
        new_comm = MPI.COMM_WORLD.Split(color = 1)

# Path: taskloaf/mpi.py
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()


# Path: tests/conftest.py
import sys
import logging
import _pytest.runner
from mpi4py import MPI
from taskloaf.mpi import MPIComm



def pytest_runtest_protocol(item, nextitem):
    n_procs = getattr(item.function, 'n_procs', 1)
    rank = MPI.COMM_WORLD.Get_rank()
    size = MPI.COMM_WORLD.Get_size()


    if n_procs > size:
        print(
            'skipping ' + str(item.name) +
            ' because it needs ' + str(n_procs) + ' MPI procs'
        )
        return True

    if rank >= n_procs:
        new_comm = MPI.COMM_WORLD.Split(color = 0)
        MPI.COMM_WORLD.Barrier()
        return True
    else:
        new_comm = MPI.COMM_WORLD.Split(color = 1)

        MPIComm.default_comm = new_comm