From 4128a40695e0e531d2f11374c4ac5cabd9d97f91 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Mon, 19 Aug 2024 14:56:48 -0700 Subject: [PATCH 1/2] add fix for from_flat --- src/nested_pandas/nestedframe/core.py | 3 ++- .../nested_pandas/nestedframe/test_nestedframe.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py index cb743fc..c7f2b38 100644 --- a/src/nested_pandas/nestedframe/core.py +++ b/src/nested_pandas/nestedframe/core.py @@ -211,7 +211,8 @@ def from_flat(cls, df, base_columns, nested_columns=None, index=None, name="nest out_df = df[base_columns][~df.index.duplicated(keep="first")] # add nested - nested_columns = [col for col in df.columns if col not in base_columns] + if nested_columns is None: + nested_columns = [col for col in df.columns if col not in base_columns] return out_df.add_nested(df[nested_columns], name=name) @classmethod diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py index 7c4d2fc..d12e32f 100644 --- a/tests/nested_pandas/nestedframe/test_nestedframe.py +++ b/tests/nested_pandas/nestedframe/test_nestedframe.py @@ -319,6 +319,20 @@ def test_recover_from_flat(): assert nf2.equals(nf) +def test_from_flat_omitting_columns(): + """test that from_flat successfully produces subsets""" + flat = NestedFrame({"a":[1,1,1,2,2], "b":[2,2,2,4,4], "c":[1,2,3,4,5], + "d":[2,4,6,8,10]}, index=[0,0,0,1,1]) + + # omit a base column + nf = NestedFrame.from_flat(flat, base_columns=["b"], nested_columns=["c", "d"]) + assert list(nf.columns) == ["b", "nested"] + assert list(nf.nested.nest.fields) == ["c", "d"] + + # omit a nested column + nf = NestedFrame.from_flat(flat, base_columns=["a", "b"], nested_columns=["c"]) + assert list(nf.columns) == ["a", "b", "nested"] + assert list(nf.nested.nest.fields) == ["c"] def test_from_lists(): """Test NestedFrame.from_lists behavior""" From 05565925d8d81e5038100d19e1c7130a7e33c8d0 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Mon, 19 Aug 2024 14:59:56 -0700 Subject: [PATCH 2/2] formatting --- tests/nested_pandas/nestedframe/test_nestedframe.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py index d12e32f..097ce1b 100644 --- a/tests/nested_pandas/nestedframe/test_nestedframe.py +++ b/tests/nested_pandas/nestedframe/test_nestedframe.py @@ -319,11 +319,14 @@ def test_recover_from_flat(): assert nf2.equals(nf) + def test_from_flat_omitting_columns(): """test that from_flat successfully produces subsets""" - flat = NestedFrame({"a":[1,1,1,2,2], "b":[2,2,2,4,4], "c":[1,2,3,4,5], - "d":[2,4,6,8,10]}, index=[0,0,0,1,1]) - + flat = NestedFrame( + {"a": [1, 1, 1, 2, 2], "b": [2, 2, 2, 4, 4], "c": [1, 2, 3, 4, 5], "d": [2, 4, 6, 8, 10]}, + index=[0, 0, 0, 1, 1], + ) + # omit a base column nf = NestedFrame.from_flat(flat, base_columns=["b"], nested_columns=["c", "d"]) assert list(nf.columns) == ["b", "nested"] @@ -334,6 +337,7 @@ def test_from_flat_omitting_columns(): assert list(nf.columns) == ["a", "b", "nested"] assert list(nf.nested.nest.fields) == ["c"] + def test_from_lists(): """Test NestedFrame.from_lists behavior""" nf = NestedFrame(