-
-
Notifications
You must be signed in to change notification settings - Fork 19.1k
ENH: Implement MultiIndex.insert_level for inserting levels at specified positions #62610
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
base: main
Are you sure you want to change the base?
Changes from all commits
e3d6970
45ac8ef
5b76304
5f0caf0
97a98e5
1a9ddc5
2199e6e
44985ad
9e8676d
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,91 @@ | ||
import pytest | ||
|
||
import pandas as pd | ||
import pandas._testing as tm | ||
|
||
|
||
class TestMultiIndexInsertLevel: | ||
def setup_method(self): | ||
self.simple_idx = pd.MultiIndex.from_tuples( | ||
[("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"] | ||
) | ||
self.empty_idx = pd.MultiIndex.from_tuples([], names=["level1", "level2"]) | ||
|
||
def test_insert_level_basic(self): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you parametrize this test? |
||
result = self.simple_idx.insert_level(0, "new_value") | ||
expected = pd.MultiIndex.from_tuples( | ||
[("new_value", "A", 1), ("new_value", "B", 2), ("new_value", "C", 3)], | ||
names=[None, "level1", "level2"], | ||
) | ||
tm.assert_index_equal(result, expected) | ||
|
||
result = self.simple_idx.insert_level(1, "middle") | ||
expected = pd.MultiIndex.from_tuples( | ||
[("A", "middle", 1), ("B", "middle", 2), ("C", "middle", 3)], | ||
names=["level1", None, "level2"], | ||
) | ||
tm.assert_index_equal(result, expected) | ||
|
||
def test_insert_level_with_different_values(self): | ||
new_values = ["X", "Y", "Z"] | ||
result = self.simple_idx.insert_level(1, new_values) | ||
expected = pd.MultiIndex.from_tuples( | ||
[("A", "X", 1), ("B", "Y", 2), ("C", "Z", 3)], | ||
names=["level1", None, "level2"], | ||
) | ||
tm.assert_index_equal(result, expected) | ||
|
||
def test_insert_level_with_name(self): | ||
result = self.simple_idx.insert_level(0, "new_val", name="new_level") | ||
assert result.names[0] == "new_level" | ||
Comment on lines
+38
to
+40
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You could also test adding names in the first test and remove this test. |
||
|
||
def test_insert_level_edge_positions(self): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This can go into the first test. |
||
result_start = self.simple_idx.insert_level(0, "start") | ||
assert result_start.nlevels == 3 | ||
|
||
result_end = self.simple_idx.insert_level(2, "end") | ||
assert result_end.nlevels == 3 | ||
|
||
def test_insert_level_error_cases(self): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you parametrize this? |
||
with pytest.raises(ValueError, match="position must be between"): | ||
self.simple_idx.insert_level(5, "invalid") | ||
|
||
with pytest.raises(ValueError, match="position must be between"): | ||
self.simple_idx.insert_level(-1, "invalid") | ||
|
||
with pytest.raises(ValueError, match="Length of values must match"): | ||
self.simple_idx.insert_level(1, ["too", "few"]) | ||
|
||
def test_insert_level_with_different_data_types(self): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These tests could go into the first test too. |
||
result_int = self.simple_idx.insert_level(1, 100) | ||
|
||
result_float = self.simple_idx.insert_level(1, 1.5) | ||
|
||
result_none = self.simple_idx.insert_level(1, None) | ||
|
||
assert result_int.nlevels == 3 | ||
assert result_float.nlevels == 3 | ||
assert result_none.nlevels == 3 | ||
|
||
def test_insert_level_preserves_original(self): | ||
original = self.simple_idx.copy() | ||
result = self.simple_idx.insert_level(1, "temp") | ||
|
||
tm.assert_index_equal(original, self.simple_idx) | ||
|
||
assert result.nlevels == original.nlevels + 1 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This assertion feels redundant against the first test. |
||
|
||
def test_debug_names(): | ||
idx = pd.MultiIndex.from_tuples( | ||
[("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"] | ||
) | ||
print("Original names:", idx.names) | ||
|
||
result = idx.insert_level(0, "new_value") | ||
print("Result names:", result.names) | ||
|
||
expected = pd.MultiIndex.from_tuples( | ||
[("new_value", "A", 1), ("new_value", "B", 2), ("new_value", "C", 3)], | ||
names=[None, "level1", "level2"], | ||
) | ||
print("Expected names:", expected.names) | ||
Comment on lines
+78
to
+91
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why does this exist? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
NIT: I would prefer that you define this in the test body that should use it.