<commit_before>import types
from functools import partial
from textwrap import dedent
from io import StringIO


import pytest
import pandas.util.testing as pdtest

import numpy
import pandas

from wqio.utils import misc
from wqio.tests import helpers


@pytest.fixture
def basic_data():
    testcsv = """\
    Date,A,B,C,D
    X,1,2,3,4
    Y,5,6,7,8
    Z,9,0,1,2
    """
    return pandas.read_csv(StringIO(dedent(testcsv)), index_col=['Date'])


@pytest.fixture
def multiindex_df():
    index = pandas.MultiIndex.from_product([['A', 'B', 'C'], ['mg/L']], names=['loc', 'units'])
    return pandas.DataFrame([[1, 2], [3, 4], [5, 6]], index=index, columns=['a', 'b'])


class mockDataset(object):
    def __init__(self, inflow, outflow):
        self.inflow = mockLocation(inflow)
        self.outflow = mockLocation(outflow)


class mockLocation(object):
    def __init__(self, data):
        self.data = data
        self.stats = mockSummary(data)


class mockSummary(object):
    def __init__(self, data):
        self.N = len(data)
        self.max = max(data)
        self.min = min(data)
        self.nonething = None


def test_addSecondColumnLevel(basic_data):
    known_cols = pandas.MultiIndex.from_tuples([(u'test', u'A'), (u'test', u'B'),
                                                (u'test', u'C'), (u'test', u'D')])
    newdata = misc.addSecondColumnLevel('test', 'testlevel', basic_data)
    assert known_cols.tolist() == newdata.columns.tolist()

    # can only add levels to non-MultiIndex columns
    with pytest.raises(ValueError):
        misc.addSecondColumnLevel('test2', 'testlevel2', newdata)


def test_add_column_level(basic_data):
    known_cols = pandas.MultiIndex.from_tuples([(u'test', u'A'), (u'test', u'B'),
                                                (u'test', u'C'), (u'test', u'D')])
    newdata = misc.add_column_level(basic_data, 'test', 'testlevel')
    assert known_cols.tolist() == newdata.columns.tolist()

    # can only add levels to non-MultiIndex columns
    with pytest.raises(ValueError):
        misc.add_column_level(newdata, 'test2', 'testlevel2')


@pytest.mark.parametrize('L1', [0, 'loc'])
@pytest.mark.parametrize('L2', [2, 'units'])
def test_swap_column_levels(multiindex_df, L1, L2):
    columns = pandas.MultiIndex.from_product(
        [['A', 'B', 'C'], ['res', 'cen'], ['mg/L']],
        names=['loc', 'value', 'units']
    )
    data = numpy.arange(len(columns) * 10).reshape((10, len(columns)))
    df = pandas.DataFrame(data, columns=columns).pipe(misc.swap_column_levels, L1, L2)

    expected_columns = pandas.MultiIndex.from_product(
        [['mg/L'], ['cen', 'res'], ['A', 'B', 'C']],
        names=['units', 'value', 'loc']
    )

    pdtest.assert_index_equal(df.columns, expected_columns)


@pytest.mark.parametrize('criteria', [None, lambda row: row[0] in ['A', 'B']])
@pytest.mark.parametrize('dropold', [True, False])
def test_redefine_index_level(multiindex_df, criteria, dropold):
    expected_cols = ['a', 'b']
    if dropold:
        expected_value = [[1, 2], [3, 4], [5, 6]]
        if criteria:
            expected_index = [('A', 'ug/L'), ('B', 'ug/L'), ('C', 'mg/L')]<commit_msg>test for misc.flatten_columns<commit_after>import types
from functools import partial
from textwrap import dedent
from io import StringIO


import pytest
import pandas.util.testing as pdtest

import numpy
import pandas

from wqio.utils import misc
from wqio.tests import helpers


@pytest.fixture
def basic_data():
    testcsv = """\
    Date,A,B,C,D
    X,1,2,3,4
    Y,5,6,7,8
    Z,9,0,1,2
    """
    return pandas.read_csv(StringIO(dedent(testcsv)), index_col=['Date'])


@pytest.fixture
def multiindex_df():
    index = pandas.MultiIndex.from_product([['A', 'B', 'C'], ['mg/L']], names=['loc', 'units'])
    return pandas.DataFrame([[1, 2], [3, 4], [5, 6]], index=index, columns=['a', 'b'])


class mockDataset(object):
    def __init__(self, inflow, outflow):
        self.inflow = mockLocation(inflow)
        self.outflow = mockLocation(outflow)


class mockLocation(object):
    def __init__(self, data):
        self.data = data
        self.stats = mockSummary(data)


class mockSummary(object):
    def __init__(self, data):
        self.N = len(data)
        self.max = max(data)
        self.min = min(data)
        self.nonething = None


def test_addSecondColumnLevel(basic_data):
    known_cols = pandas.MultiIndex.from_tuples([(u'test', u'A'), (u'test', u'B'),
                                                (u'test', u'C'), (u'test', u'D')])
    newdata = misc.addSecondColumnLevel('test', 'testlevel', basic_data)
    assert known_cols.tolist() == newdata.columns.tolist()

    # can only add levels to non-MultiIndex columns
    with pytest.raises(ValueError):
        misc.addSecondColumnLevel('test2', 'testlevel2', newdata)


def test_add_column_level(basic_data):
    known_cols = pandas.MultiIndex.from_tuples([(u'test', u'A'), (u'test', u'B'),
                                                (u'test', u'C'), (u'test', u'D')])
    newdata = misc.add_column_level(basic_data, 'test', 'testlevel')
    assert known_cols.tolist() == newdata.columns.tolist()

    # can only add levels to non-MultiIndex columns
    with pytest.raises(ValueError):
        misc.add_column_level(newdata, 'test2', 'testlevel2')


@pytest.mark.parametrize('L1', [0, 'loc'])
@pytest.mark.parametrize('L2', [2, 'units'])
def test_swap_column_levels(multiindex_df, L1, L2):
    columns = pandas.MultiIndex.from_product(
        [['A', 'B', 'C'], ['res', 'cen'], ['mg/L']],
        names=['loc', 'value', 'units']
    )
    data = numpy.arange(len(columns) * 10).reshape((10, len(columns)))
    df = pandas.DataFrame(data, columns=columns).pipe(misc.swap_column_levels, L1, L2)

    expected_columns = pandas.MultiIndex.from_product(
        [['mg/L'], ['cen', 'res'], ['A', 'B', 'C']],
        names=['units', 'value', 'loc']
    )

    pdtest.assert_index_equal(df.columns, expected_columns)


@pytest.mark.parametrize(('df', 'expected'), [
    (multiindex_df().T, ['A_mg/L', 'B_mg/L', 'C_mg/L']),
    (basic_data(), ['A', 'B', 'C', 'D']),
])
def test_flatten_columns(df, expected):
    flat = misc.flatten_columns(df)
    assert flat.columns.tolist() == expected


@pytest.mark.parametrize('criteria', [None, lambda row: row[0] in ['A', 'B']])
@pytest.mark.parametrize('dropold', [True, False])
def test_redefine_index_level(multiindex_df, criteria, dropold):
    expected_cols = ['a', 'b']
    if dropold:
        expected_value = [[1, 2], [3, 4], [5, 6]]
        if criteria:
            expected_index = [('A', 'ug/L'), ('B', 'ug/L'), ('C', 'mg/L')]
{'fraction_non_alphanumeric': 0.1770298822819197, 'fraction_numerical': 0.012375490492001207, 'mean_word_length': 3.4898373983739837, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4651598', 'n_tokens_mistral': 2408, 'n_tokens_neox': 2341, 'n_words': 488}