|
1 |
| -from pandas import DataFrame |
| 1 | +import pandas as pd |
2 | 2 |
|
3 | 3 |
|
4 |
| -class CustomDataFrame(DataFrame): |
| 4 | +class SubclassedDataFrame2(pd.DataFrame): |
5 | 5 | """
|
6 | 6 | Extension of DataFrame as described in [guidelines]
|
7 | 7 |
|
8 | 8 | [guidelines]: https://pandas.pydata.org/pandas-docs/stable/development/extending.html#override-constructor-properties # noqa
|
9 | 9 | """
|
10 | 10 |
|
11 |
| - _metadata = ["_custom_metadata"] |
| 11 | + # temporary properties |
| 12 | + _internal_names = pd.DataFrame._internal_names + ["internal_cache"] |
| 13 | + _internal_names_set = set(_internal_names) |
| 14 | + |
| 15 | + # normal properties |
| 16 | + _metadata = ["added_property"] |
12 | 17 |
|
13 | 18 | @property
|
14 | 19 | def _constructor(self):
|
15 |
| - return CustomDataFrame |
| 20 | + return SubclassedDataFrame2 |
16 | 21 |
|
17 | 22 |
|
18 | 23 | def test_groupby_with_custom_metadata():
|
19 |
| - custom_df = CustomDataFrame( |
| 24 | + custom_df = SubclassedDataFrame2( |
20 | 25 | [[11, 12, 0], [21, 22, 0], [31, 32, 1]], columns=["a", "b", "g"]
|
21 | 26 | )
|
22 |
| - custom_df._custom_metadata = "Custom metadata" |
23 |
| - |
24 |
| - for _, group_df in custom_df.groupby("g"): |
25 |
| - assert group_df._custom_metadata == "Custom metadata" |
| 27 | + custom_df.added_property = "hello_pandas" |
| 28 | + grouped = custom_df.groupby("g") |
| 29 | + for _, group_df in grouped: |
| 30 | + assert group_df.added_property == "hello_pandas" |
| 31 | + |
| 32 | + |
| 33 | +def test_groupby_sum_with_custom_metadata(): |
| 34 | + my_data_as_dictionary = { |
| 35 | + "mycategorical": [1, 1, 2], |
| 36 | + "myfloat1": [1.0, 2.0, 3.0], |
| 37 | + "myfloat2": [1.0, 2.0, 3.0], |
| 38 | + } |
| 39 | + sdf = SubclassedDataFrame2(my_data_as_dictionary) |
| 40 | + sdf.added_property = "hello pandas" |
| 41 | + grouped = sdf.groupby("mycategorical")[["myfloat1", "myfloat2"]] |
| 42 | + df = grouped.sum() |
| 43 | + assert df.added_property == "hello pandas" |
0 commit comments