Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.16.1.txt
Original file line number Diff line number Diff line change
Expand Up @@ -300,3 +300,4 @@ Bug Fixes
- Bug in ``transform`` when groups are equal in number and dtype to the input index (:issue:`9700`)
- Google BigQuery connector now imports dependencies on a per-method basis.(:issue:`9713`)
- Updated BigQuery connector to no longer use deprecated ``oauth2client.tools.run()`` (:issue:`8327`)
- Bug in subclassed ``DataFrame``. It may not return the correct class, when slicing or subsetting it. (:issue:`9632`)
4 changes: 2 additions & 2 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -1839,15 +1839,15 @@ def _getitem_multilevel(self, key):
result.columns = result_columns
else:
new_values = self.values[:, loc]
result = DataFrame(new_values, index=self.index,
result = self._constructor(new_values, index=self.index,
columns=result_columns).__finalize__(self)
if len(result.columns) == 1:
top = result.columns[0]
if ((type(top) == str and top == '') or
(type(top) == tuple and top[0] == '')):
result = result['']
if isinstance(result, Series):
result = Series(result, index=self.index, name=key)
result = self._constructor_sliced(result, index=self.index, name=key)

result._set_is_copy(self)
return result
Expand Down
53 changes: 53 additions & 0 deletions pandas/tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -2791,6 +2791,59 @@ def test_insert_error_msmgs(self):
with assertRaisesRegexp(TypeError, msg):
df['gr'] = df.groupby(['b', 'c']).count()

def test_frame_subclassing_and_slicing(self):
# Subclass frame and ensure it returns the right class on slicing it
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add the issue (this pr number is ok) as a comment

# In reference to PR 9632

class CustomSeries(Series):
@property
def _constructor(self):
return CustomSeries

def custom_series_function(self):
return 'OK'

class CustomDataFrame(DataFrame):
"Subclasses pandas DF, fills DF with simulation results, adds some custom plotting functions."

def __init__(self, *args, **kw):
super(CustomDataFrame, self).__init__(*args, **kw)

@property
def _constructor(self):
return CustomDataFrame

_constructor_sliced = CustomSeries

def custom_frame_function(self):
return 'OK'

data = {'col1': range(10),
'col2': range(10)}
cdf = CustomDataFrame(data)

# Did we get back our own DF class?
self.assertTrue(isinstance(cdf, CustomDataFrame))

# Do we get back our own Series class after selecting a column?
cdf_series = cdf.col1
self.assertTrue(isinstance(cdf_series, CustomSeries))
self.assertEqual(cdf_series.custom_series_function(), 'OK')

# Do we get back our own DF class after slicing row-wise?
cdf_rows = cdf[1:5]
self.assertTrue(isinstance(cdf_rows, CustomDataFrame))
self.assertEqual(cdf_rows.custom_frame_function(), 'OK')

# Make sure sliced part of multi-index frame is custom class
mcol = pd.MultiIndex.from_tuples([('A', 'A'), ('A', 'B')])
cdf_multi = CustomDataFrame([[0, 1], [2, 3]], columns=mcol)
self.assertTrue(isinstance(cdf_multi['A'], CustomDataFrame))

mcol = pd.MultiIndex.from_tuples([('A', ''), ('B', '')])
cdf_multi2 = CustomDataFrame([[0, 1], [2, 3]], columns=mcol)
self.assertTrue(isinstance(cdf_multi2['A'], CustomSeries))

def test_constructor_subclass_dict(self):
# Test for passing dict subclass to constructor
data = {'col1': tm.TestSubDict((x, 10.0 * x) for x in range(10)),
Expand Down