Repository: lincc-frameworks/nested-pandas

Issue:
NestedFrame.from_flat doesn't work with pd.DataFrame input
**Bug report**

```python
NestedFrame.from_flat(
    pd.DataFrame({'a': [1, 2], 'b': [3, None]}, index=[1, 1]),
    base_columns=[],
)
```

```
AttributeError                            Traceback (most recent call last)
<ipython-input-5-4df41e18d208> in ?()
----> 1 NestedFrame.from_flat(pd.DataFrame({'a': [1, 2], 'b': [3, None]}, index=[1,1]), base_columns=[])

~/.virtualenvs/v/lib/python3.12/site-packages/nested_pandas/nestedframe/core.py in ?(cls, df, base_columns, nested_columns, index, name)
    212
    213         # add nested
    214         if nested_columns is None:
    215             nested_columns = [col for col in df.columns if col not in base_columns]
--> 216         return out_df.add_nested(df[nested_columns], name=name)

~/.virtualenvs/v/lib/python3.12/site-packages/pandas/core/generic.py in ?(self, name)
   6295             and name not in self._accessors
   6296             and self._info_axis._can_hold_identifiers_and_holds_name(name)
   6297         ):
   6298             return self[name]
-> 6299         return object.__getattribute__(self, name)

AttributeError: 'DataFrame' object has no attribute 'add_nested'
```

We should probably wrap out_df with `NestedFrame` here:
https://github.com/lincc-frameworks/nested-pandas/blob/6101d396bc9f3d281393ff3bd6da819d6a473ded/src/nested_pandas/nestedframe/core.py#L216

**Before submitting**
Please check the following:

- [x] I have described the situation in which the bug arose, including what code was executed, information about my environment, and any applicable data others will need to reproduce the problem.
- [x] I have included available evidence of the unexpected behavior (including error messages, screenshots, and/or plots) as well as a descriprion of what I expected instead.
- [x] If I have a solution in mind, I have provided an explanation and/or pseudocode and/or task list.

Discussion:

Patch:
diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py
index c7f2b38..3f28a5e 100644
--- a/src/nested_pandas/nestedframe/core.py
+++ b/src/nested_pandas/nestedframe/core.py
@@ -210,6 +210,10 @@ class NestedFrame(pd.DataFrame):
         # drop duplicates on index
         out_df = df[base_columns][~df.index.duplicated(keep="first")]
 
+        # Convert df to NestedFrame if needed
+        if not isinstance(out_df, NestedFrame):
+            out_df = NestedFrame(out_df)
+
         # add nested
         if nested_columns is None:
             nested_columns = [col for col in df.columns if col not in base_columns]


Test patch:
diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py
index 097ce1b..95421b3 100644
--- a/tests/nested_pandas/nestedframe/test_nestedframe.py
+++ b/tests/nested_pandas/nestedframe/test_nestedframe.py
@@ -285,13 +285,21 @@ def test_add_nested_for_empty_df():
     assert_frame_equal(new_base.nested.nest.to_flat(), nested.astype(pd.ArrowDtype(pa.float64())))
 
 
+@pytest.mark.parametrize("pandas", [False, True])
 @pytest.mark.parametrize("index", [None, "a", "c"])
-def test_from_flat(index):
+def test_from_flat(index, pandas):
     """Test the NestedFrame.from_flat functionality"""
-    nf = 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],
-    )
+
+    if pandas:
+        nf = pd.DataFrame(
+            {"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],
+        )
+    else:
+        nf = 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],
+        )
 
     out_nf = NestedFrame.from_flat(nf, base_columns=["a", "b"], index=index, name="new_nested")