<commit_before>from ..tester import Tester as BaseTester

import pytest
import traceback
import warnings
import sys
import os
import contextlib
import time
from argparse import ArgumentParser
from contextlib import contextmanager

if sys.version_info[0] == 2:
    from cStringIO import StringIO
else:
    from io import StringIO
import re

def fix_titles(s):
    pattern = '=====+'
    return re.sub(pattern, lambda x: x.group(0).replace('=', '-'), s)

class Rotator(object):
    """ in a rotator every range runs in terms """
    def __init__(self, comm):
        self.comm = comm
    def __enter__(self):
        self.comm.Barrier()
        for i in range(self.comm.rank):
            self.comm.Barrier()
    def __exit__(self, type, value, tb):
        for i in range(self.comm.rank, self.comm.size):
            self.comm.Barrier()
        self.comm.Barrier()

@contextmanager
def nompi(comm):
    errored = False
    error = None
    try:
        yield
    except Exception as e:
        errored = True
        error = e
    finally:
        anyerrored = any(comm.allgather(errored))

    if anyerrored:
        if error is None:
            raise RuntimeError("Some ranks failed")
        else:
            raise error


communicators = {
}

class WorldTooSmall(Exception): pass

def create_comm(size):
    from mpi4py import MPI
    if MPI.COMM_WORLD.size < size:
        raise WorldTooSmall

    color = 0 if MPI.COMM_WORLD.rank < size else 1
    if size not in communicators:
        if MPI.COMM_WORLD.size == size:
            comm = MPI.COMM_WORLD
        elif size == 1:
            comm = MPI.COMM_SELF
        else:
            comm = MPI.COMM_WORLD.Split(color)
        communicators[size] = comm

    return communicators[size], color

def MPITestFixture(commsize, scope='function'):
    """ Create a test fixture for MPI Communicators of various commsizes """

    @pytest.fixture(params=commsize, scope=scope)
    def fixture(request):
        from mpi4py import MPI
        MPI.COMM_WORLD.barrier()
        try:
            comm, color = create_comm(request.param)

            # FIXME: use skip after https://github.com/pytest-dev/pytest/issues/3969 is figured out.
            if color != 0:
                request.applymarker(pytest.mark.xfail(reason="Not using communicator %d" % request.param, run=False))
                return None
            else:
                return comm

        except WorldTooSmall:
            # FIXME: use skip after https://github.com/pytest-dev/pytest/issues/3969 is figured out.
            request.applymarker(pytest.mark.xfail(reason="World is too small", run=False))
            return None

    return fixture

def MPITest(commsize):
    """
    A decorator that repeatedly calls the wrapped function,
    with communicators of varying sizes.

    This converts the test to a generator test; therefore the
    underlyig test shall not be a generator test.<commit_msg>use skip after https://github.com/pytest-dev/pytest/issues/3969 instead of xfail<commit_after>from ..tester import Tester as BaseTester

import pytest
import traceback
import warnings
import sys
import os
import contextlib
import time
from argparse import ArgumentParser
from contextlib import contextmanager

if sys.version_info[0] == 2:
    from cStringIO import StringIO
else:
    from io import StringIO
import re

def fix_titles(s):
    pattern = '=====+'
    return re.sub(pattern, lambda x: x.group(0).replace('=', '-'), s)

class Rotator(object):
    """ in a rotator every range runs in terms """
    def __init__(self, comm):
        self.comm = comm
    def __enter__(self):
        self.comm.Barrier()
        for i in range(self.comm.rank):
            self.comm.Barrier()
    def __exit__(self, type, value, tb):
        for i in range(self.comm.rank, self.comm.size):
            self.comm.Barrier()
        self.comm.Barrier()

@contextmanager
def nompi(comm):
    errored = False
    error = None
    try:
        yield
    except Exception as e:
        errored = True
        error = e
    finally:
        anyerrored = any(comm.allgather(errored))

    if anyerrored:
        if error is None:
            raise RuntimeError("Some ranks failed")
        else:
            raise error


communicators = {
}

class WorldTooSmall(Exception): pass

def create_comm(size):
    from mpi4py import MPI
    if MPI.COMM_WORLD.size < size:
        raise WorldTooSmall

    color = 0 if MPI.COMM_WORLD.rank < size else 1
    if size not in communicators:
        if MPI.COMM_WORLD.size == size:
            comm = MPI.COMM_WORLD
        elif size == 1:
            comm = MPI.COMM_SELF
        else:
            comm = MPI.COMM_WORLD.Split(color)
        communicators[size] = comm

    return communicators[size], color

def MPITestFixture(commsize, scope='function'):
    """ Create a test fixture for MPI Communicators of various commsizes """

    @pytest.fixture(params=commsize, scope=scope)
    def fixture(request):
        from mpi4py import MPI
        MPI.COMM_WORLD.barrier()
        try:
            comm, color = create_comm(request.param)

            if color != 0:
                pytest.skip("Not using communicator %d" %(request.param))
                return None
            else:
                return comm

        except WorldTooSmall:
            pytest.skip("Not using communicator %d" % request.param)
            return None

    return fixture

def MPITest(commsize):
    """
    A decorator that repeatedly calls the wrapped function,
    with communicators of varying sizes.

    This converts the test to a generator test; therefore the
    underlyig test shall not be a generator test.
{'fraction_non_alphanumeric': 0.08803252607389075, 'fraction_numerical': 0.00530316422131872, 'mean_word_length': 2.3618538324420677, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2721874', 'n_tokens_mistral': 1720, 'n_tokens_neox': 1651, 'n_words': 516}