Skip to content

REF/TST: Add more pytest idiom to indexing/multiindex/test_iloc.py #24272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Dec 14, 2018
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
232 changes: 110 additions & 122 deletions pandas/tests/indexing/multiindex/test_iloc.py
Original file line number Diff line number Diff line change
@@ -1,129 +1,117 @@
from warnings import catch_warnings

import numpy as np
import pytest

from pandas import DataFrame, MultiIndex, Series
from pandas.util import testing as tm


@pytest.mark.filterwarnings("ignore:\\n.ix:DeprecationWarning")
class TestMultiIndexIloc(object):

def test_iloc_getitem_multiindex2(self):
# TODO(wesm): fix this
pytest.skip('this test was being suppressed, '
'needs to be fixed')

arr = np.random.randn(3, 3)
df = DataFrame(arr, columns=[[2, 2, 4], [6, 8, 10]],
index=[[4, 4, 8], [8, 10, 12]])

rs = df.iloc[2]
xp = Series(arr[2], index=df.columns)
tm.assert_series_equal(rs, xp)

rs = df.iloc[:, 2]
xp = Series(arr[:, 2], index=df.index)
tm.assert_series_equal(rs, xp)

rs = df.iloc[2, 2]
xp = df.values[2, 2]
assert rs == xp

# for multiple items
# GH 5528
rs = df.iloc[[0, 1]]
xp = df.xs(4, drop_level=False)
tm.assert_frame_equal(rs, xp)

tup = zip(*[['a', 'a', 'b', 'b'], ['x', 'y', 'x', 'y']])
index = MultiIndex.from_tuples(tup)
df = DataFrame(np.random.randn(4, 4), index=index)
rs = df.iloc[[2, 3]]
xp = df.xs('b', drop_level=False)
tm.assert_frame_equal(rs, xp)

def test_iloc_getitem_multiindex(self):
mi_labels = DataFrame(np.random.randn(4, 3),
columns=[['i', 'i', 'j'], ['A', 'A', 'B']],
index=[['i', 'i', 'j', 'k'],
['X', 'X', 'Y', 'Y']])

mi_int = DataFrame(np.random.randn(3, 3),
columns=[[2, 2, 4], [6, 8, 10]],
index=[[4, 4, 8], [8, 10, 12]])

# the first row
rs = mi_int.iloc[0]
with catch_warnings(record=True):
xp = mi_int.ix[4].ix[8]
tm.assert_series_equal(rs, xp, check_names=False)
assert rs.name == (4, 8)
assert xp.name == 8

# 2nd (last) columns
rs = mi_int.iloc[:, 2]
with catch_warnings(record=True):
xp = mi_int.ix[:, 2]
tm.assert_series_equal(rs, xp)

# corner column
rs = mi_int.iloc[2, 2]
with catch_warnings(record=True):
# First level is int - so use .loc rather than .ix (GH 21593)
xp = mi_int.loc[(8, 12), (4, 10)]
assert rs == xp

# this is basically regular indexing
rs = mi_labels.iloc[2, 2]
with catch_warnings(record=True):
xp = mi_labels.ix['j'].ix[:, 'j'].ix[0, 0]
assert rs == xp

def test_frame_getitem_setitem_slice(
self, multiindex_dataframe_random_data):
frame = multiindex_dataframe_random_data
# getitem
result = frame.iloc[:4]
expected = frame[:4]
tm.assert_frame_equal(result, expected)

# setitem
cp = frame.copy()
cp.iloc[:4] = 0

assert (cp.values[:4] == 0).all()
assert (cp.values[4:] != 0).all()

def test_indexing_ambiguity_bug_1678(self):
columns = MultiIndex.from_tuples([('Ohio', 'Green'), ('Ohio', 'Red'), (
'Colorado', 'Green')])
index = MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1), ('b', 2)
])

frame = DataFrame(np.arange(12).reshape((4, 3)), index=index,
columns=columns)

result = frame.iloc[:, 1]
exp = frame.loc[:, ('Ohio', 'Red')]
assert isinstance(result, Series)
tm.assert_series_equal(result, exp)

def test_iloc_mi(self):
# GH 13797
# Test if iloc can handle integer locations in MultiIndexed DataFrame

data = [['str00', 'str01'], ['str10', 'str11'], ['str20', 'srt21'],
['str30', 'str31'], ['str40', 'str41']]

mi = MultiIndex.from_tuples(
[('CC', 'A'), ('CC', 'B'), ('CC', 'B'), ('BB', 'a'), ('BB', 'b')])

expected = DataFrame(data)
df_mi = DataFrame(data, index=mi)

result = DataFrame([[df_mi.iloc[r, c] for c in range(2)]
for r in range(5)])

tm.assert_frame_equal(result, expected)
@pytest.fixture
def simple_multiindex_dataframe():
"""
Factory function to create simple 3 x 3 dataframe with
both columns and row MultiIndex using supplied data or
random data by default.
"""
def _simple_multiindex_dataframe(data=None):
if data is None:
data = np.random.randn(3, 3)
return DataFrame(data, columns=[[2, 2, 4], [6, 8, 10]],
index=[[4, 4, 8], [8, 10, 12]])
return _simple_multiindex_dataframe


@pytest.mark.parametrize('indexer, expected', [
(lambda df: df.iloc[0],
lambda arr: Series(arr[0], index=[[2, 2, 4], [6, 8, 10]], name=(4, 8))),
(lambda df: df.iloc[2],
lambda arr: Series(arr[2], index=[[2, 2, 4], [6, 8, 10]], name=(8, 12))),
(lambda df: df.iloc[:, 2],
lambda arr: Series(
arr[:, 2], index=[[4, 4, 8], [8, 10, 12]], name=(4, 10)))
])
def test_iloc_returns_series(indexer, expected, simple_multiindex_dataframe):
arr = np.random.randn(3, 3)
df = simple_multiindex_dataframe(arr)
result = indexer(df)
expected = expected(arr)
tm.assert_series_equal(result, expected)


def test_iloc_returns_dataframe(simple_multiindex_dataframe):
df = simple_multiindex_dataframe()
result = df.iloc[[0, 1]]
expected = df.xs(4, drop_level=False)
tm.assert_frame_equal(result, expected)


def test_iloc_returns_scalar(simple_multiindex_dataframe):
arr = np.random.randn(3, 3)
df = simple_multiindex_dataframe(arr)
result = df.iloc[2, 2]
expected = arr[2, 2]
assert result == expected


def test_iloc_getitem_multiple_items():
# GH 5528
tup = zip(*[['a', 'a', 'b', 'b'], ['x', 'y', 'x', 'y']])
index = MultiIndex.from_tuples(tup)
df = DataFrame(np.random.randn(4, 4), index=index)
result = df.iloc[[2, 3]]
expected = df.xs('b', drop_level=False)
tm.assert_frame_equal(result, expected)


def test_iloc_getitem_labels():
# this is basically regular indexing
arr = np.random.randn(4, 3)
df = DataFrame(arr,
columns=[['i', 'i', 'j'], ['A', 'A', 'B']],
index=[['i', 'i', 'j', 'k'], ['X', 'X', 'Y', 'Y']])
result = df.iloc[2, 2]
expected = arr[2, 2]
assert result == expected


def test_frame_getitem_slice(multiindex_dataframe_random_data):
df = multiindex_dataframe_random_data
result = df.iloc[:4]
expected = df[:4]
tm.assert_frame_equal(result, expected)


def test_frame_setitem_slice(multiindex_dataframe_random_data):
df = multiindex_dataframe_random_data
df.iloc[:4] = 0

assert (df.values[:4] == 0).all()
assert (df.values[4:] != 0).all()


def test_indexing_ambiguity_bug_1678():
# GH 1678
columns = MultiIndex.from_tuples(
[('Ohio', 'Green'), ('Ohio', 'Red'), ('Colorado', 'Green')])
index = MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1), ('b', 2)])

df = DataFrame(np.arange(12).reshape((4, 3)), index=index, columns=columns)

result = df.iloc[:, 1]
expected = df.loc[:, ('Ohio', 'Red')]
tm.assert_series_equal(result, expected)


def test_iloc_integer_locations():
# GH 13797
data = [['str00', 'str01'], ['str10', 'str11'], ['str20', 'srt21'],
['str30', 'str31'], ['str40', 'str41']]

index = MultiIndex.from_tuples(
[('CC', 'A'), ('CC', 'B'), ('CC', 'B'), ('BB', 'a'), ('BB', 'b')])

expected = DataFrame(data)
df = DataFrame(data, index=index)

result = DataFrame([[df.iloc[r, c] for c in range(2)] for r in range(5)])

tm.assert_frame_equal(result, expected)