Skip to content

Commit

Permalink
CLN: Use ABCs in set_index
Browse files Browse the repository at this point in the history
  • Loading branch information
h-vetinari committed Feb 3, 2019
1 parent d3c9d6e commit 89ed35d
Showing 1 changed file with 10 additions and 16 deletions.
26 changes: 10 additions & 16 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4166,37 +4166,31 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
names = []
if append:
names = [x for x in self.index.names]
if isinstance(self.index, MultiIndex):
if isinstance(self.index, ABCMultiIndex):
for i in range(self.index.nlevels):
arrays.append(self.index._get_level_values(i))
else:
arrays.append(self.index)

to_remove = []
for col in keys:
if isinstance(col, MultiIndex):
# append all but the last column so we don't have to modify
# the end of this loop
for n in range(col.nlevels - 1):
if isinstance(col, ABCMultiIndex):
for n in range(col.nlevels):
arrays.append(col._get_level_values(n))

level = col._get_level_values(col.nlevels - 1)
names.extend(col.names)
elif isinstance(col, Series):
level = col._values
names.append(col.name)
elif isinstance(col, Index):
level = col
elif isinstance(col, (ABCIndexClass, ABCSeries)):
# if Index then not MultiIndex (treated above)
arrays.append(col)
names.append(col.name)
elif isinstance(col, (list, np.ndarray, Index)):
level = col
elif isinstance(col, (list, np.ndarray)):
arrays.append(col)
names.append(None)
# from here, col can only be a column label
else:
level = frame[col]._values
arrays.append(frame[col]._values)
names.append(col)
if drop:
to_remove.append(col)
arrays.append(level)

index = ensure_index_from_sequences(arrays, names)

Expand Down

0 comments on commit 89ed35d

Please sign in to comment.