261 lines
6.5 KiB
Python
261 lines
6.5 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
import pandas as pd
|
|
from pandas import (
|
|
Index,
|
|
MultiIndex,
|
|
date_range,
|
|
period_range,
|
|
)
|
|
import pandas._testing as tm
|
|
from pandas.core.api import UInt64Index
|
|
|
|
|
|
def test_shift(idx):
|
|
|
|
# GH8083 test the base class for shift
|
|
msg = "This method is only implemented for DatetimeIndex, PeriodIndex and "
|
|
"TimedeltaIndex; Got type MultiIndex"
|
|
with pytest.raises(NotImplementedError, match=msg):
|
|
idx.shift(1)
|
|
with pytest.raises(NotImplementedError, match=msg):
|
|
idx.shift(1, 2)
|
|
|
|
|
|
def test_groupby(idx):
|
|
groups = idx.groupby(np.array([1, 1, 1, 2, 2, 2]))
|
|
labels = idx.tolist()
|
|
exp = {1: labels[:3], 2: labels[3:]}
|
|
tm.assert_dict_equal(groups, exp)
|
|
|
|
# GH5620
|
|
groups = idx.groupby(idx)
|
|
exp = {key: [key] for key in idx}
|
|
tm.assert_dict_equal(groups, exp)
|
|
|
|
|
|
def test_truncate_multiindex():
|
|
# GH 34564 for MultiIndex level names check
|
|
major_axis = Index(list(range(4)))
|
|
minor_axis = Index(list(range(2)))
|
|
|
|
major_codes = np.array([0, 0, 1, 2, 3, 3])
|
|
minor_codes = np.array([0, 1, 0, 1, 0, 1])
|
|
|
|
index = MultiIndex(
|
|
levels=[major_axis, minor_axis],
|
|
codes=[major_codes, minor_codes],
|
|
names=["L1", "L2"],
|
|
)
|
|
|
|
result = index.truncate(before=1)
|
|
assert "foo" not in result.levels[0]
|
|
assert 1 in result.levels[0]
|
|
assert index.names == result.names
|
|
|
|
result = index.truncate(after=1)
|
|
assert 2 not in result.levels[0]
|
|
assert 1 in result.levels[0]
|
|
assert index.names == result.names
|
|
|
|
result = index.truncate(before=1, after=2)
|
|
assert len(result.levels[0]) == 2
|
|
assert index.names == result.names
|
|
|
|
msg = "after < before"
|
|
with pytest.raises(ValueError, match=msg):
|
|
index.truncate(3, 1)
|
|
|
|
|
|
# TODO: reshape
|
|
|
|
|
|
def test_reorder_levels(idx):
|
|
# this blows up
|
|
with pytest.raises(IndexError, match="^Too many levels"):
|
|
idx.reorder_levels([2, 1, 0])
|
|
|
|
|
|
def test_numpy_repeat():
|
|
reps = 2
|
|
numbers = [1, 2, 3]
|
|
names = np.array(["foo", "bar"])
|
|
|
|
m = MultiIndex.from_product([numbers, names], names=names)
|
|
expected = MultiIndex.from_product([numbers, names.repeat(reps)], names=names)
|
|
tm.assert_index_equal(np.repeat(m, reps), expected)
|
|
|
|
msg = "the 'axis' parameter is not supported"
|
|
with pytest.raises(ValueError, match=msg):
|
|
np.repeat(m, reps, axis=1)
|
|
|
|
|
|
def test_append_mixed_dtypes():
|
|
# GH 13660
|
|
dti = date_range("2011-01-01", freq="M", periods=3)
|
|
dti_tz = date_range("2011-01-01", freq="M", periods=3, tz="US/Eastern")
|
|
pi = period_range("2011-01", freq="M", periods=3)
|
|
|
|
mi = MultiIndex.from_arrays(
|
|
[[1, 2, 3], [1.1, np.nan, 3.3], ["a", "b", "c"], dti, dti_tz, pi]
|
|
)
|
|
assert mi.nlevels == 6
|
|
|
|
res = mi.append(mi)
|
|
exp = MultiIndex.from_arrays(
|
|
[
|
|
[1, 2, 3, 1, 2, 3],
|
|
[1.1, np.nan, 3.3, 1.1, np.nan, 3.3],
|
|
["a", "b", "c", "a", "b", "c"],
|
|
dti.append(dti),
|
|
dti_tz.append(dti_tz),
|
|
pi.append(pi),
|
|
]
|
|
)
|
|
tm.assert_index_equal(res, exp)
|
|
|
|
other = MultiIndex.from_arrays(
|
|
[
|
|
["x", "y", "z"],
|
|
["x", "y", "z"],
|
|
["x", "y", "z"],
|
|
["x", "y", "z"],
|
|
["x", "y", "z"],
|
|
["x", "y", "z"],
|
|
]
|
|
)
|
|
|
|
res = mi.append(other)
|
|
exp = MultiIndex.from_arrays(
|
|
[
|
|
[1, 2, 3, "x", "y", "z"],
|
|
[1.1, np.nan, 3.3, "x", "y", "z"],
|
|
["a", "b", "c", "x", "y", "z"],
|
|
dti.append(Index(["x", "y", "z"])),
|
|
dti_tz.append(Index(["x", "y", "z"])),
|
|
pi.append(Index(["x", "y", "z"])),
|
|
]
|
|
)
|
|
tm.assert_index_equal(res, exp)
|
|
|
|
|
|
def test_iter(idx):
|
|
result = list(idx)
|
|
expected = [
|
|
("foo", "one"),
|
|
("foo", "two"),
|
|
("bar", "one"),
|
|
("baz", "two"),
|
|
("qux", "one"),
|
|
("qux", "two"),
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
def test_sub(idx):
|
|
|
|
first = idx
|
|
|
|
# - now raises (previously was set op difference)
|
|
msg = "cannot perform __sub__ with this index type: MultiIndex"
|
|
with pytest.raises(TypeError, match=msg):
|
|
first - idx[-3:]
|
|
with pytest.raises(TypeError, match=msg):
|
|
idx[-3:] - first
|
|
with pytest.raises(TypeError, match=msg):
|
|
idx[-3:] - first.tolist()
|
|
msg = "cannot perform __rsub__ with this index type: MultiIndex"
|
|
with pytest.raises(TypeError, match=msg):
|
|
first.tolist() - idx[-3:]
|
|
|
|
|
|
def test_map(idx):
|
|
# callable
|
|
index = idx
|
|
|
|
result = index.map(lambda x: x)
|
|
tm.assert_index_equal(result, index)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"mapper",
|
|
[
|
|
lambda values, idx: {i: e for e, i in zip(values, idx)},
|
|
lambda values, idx: pd.Series(values, idx),
|
|
],
|
|
)
|
|
def test_map_dictlike(idx, mapper):
|
|
|
|
identity = mapper(idx.values, idx)
|
|
|
|
# we don't infer to UInt64 for a dict
|
|
if isinstance(idx, UInt64Index) and isinstance(identity, dict):
|
|
expected = idx.astype("int64")
|
|
else:
|
|
expected = idx
|
|
|
|
result = idx.map(identity)
|
|
tm.assert_index_equal(result, expected)
|
|
|
|
# empty mappable
|
|
expected = Index([np.nan] * len(idx))
|
|
result = idx.map(mapper(expected, idx))
|
|
tm.assert_index_equal(result, expected)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"func",
|
|
[
|
|
np.exp,
|
|
np.exp2,
|
|
np.expm1,
|
|
np.log,
|
|
np.log2,
|
|
np.log10,
|
|
np.log1p,
|
|
np.sqrt,
|
|
np.sin,
|
|
np.cos,
|
|
np.tan,
|
|
np.arcsin,
|
|
np.arccos,
|
|
np.arctan,
|
|
np.sinh,
|
|
np.cosh,
|
|
np.tanh,
|
|
np.arcsinh,
|
|
np.arccosh,
|
|
np.arctanh,
|
|
np.deg2rad,
|
|
np.rad2deg,
|
|
],
|
|
ids=lambda func: func.__name__,
|
|
)
|
|
def test_numpy_ufuncs(idx, func):
|
|
# test ufuncs of numpy. see:
|
|
# https://numpy.org/doc/stable/reference/ufuncs.html
|
|
|
|
expected_exception = TypeError
|
|
msg = (
|
|
"loop of ufunc does not support argument 0 of type tuple which "
|
|
f"has no callable {func.__name__} method"
|
|
)
|
|
with pytest.raises(expected_exception, match=msg):
|
|
func(idx)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"func",
|
|
[np.isfinite, np.isinf, np.isnan, np.signbit],
|
|
ids=lambda func: func.__name__,
|
|
)
|
|
def test_numpy_type_funcs(idx, func):
|
|
msg = (
|
|
f"ufunc '{func.__name__}' not supported for the input types, and the inputs "
|
|
"could not be safely coerced to any supported types according to "
|
|
"the casting rule ''safe''"
|
|
)
|
|
with pytest.raises(TypeError, match=msg):
|
|
func(idx)
|