ENH: pass single label to drop. GH #1143

Author: Chang She

pandas/core/common.py

@@ -608,6 +608,20 @@ def _asarray_tuplesafe(values, dtype=None):
 
     return result
 
+def _index_labels_to_array(labels):
+    if isinstance(labels, (basestring, tuple)):
+        labels = [labels]
+
+    if not isinstance(labels, (list, np.ndarray)):
+        try:
+            labels = list(labels)
+        except TypeError: # non-iterable
+            labels = [labels]
+
+    labels = _asarray_tuplesafe(labels)
+
+    return labels
+
 def _stringify(col):
     # unicode workaround
     try:

pandas/core/index.py

@@ -1011,7 +1011,7 @@ def drop(self, labels):
         -------
         dropped : Index
         """
-        labels = np.asarray(list(labels), dtype=object)
+        labels = com._index_labels_to_array(labels)
         indexer = self.get_indexer(labels)
         mask = indexer == -1
         if mask.any():
@@ -1511,7 +1511,7 @@ def drop(self, labels, level=None):
 
         try:
             if not isinstance(labels, np.ndarray):
-                labels = com._asarray_tuplesafe(labels)
+                labels = com._index_labels_to_array(labels)
             indexer = self.get_indexer(labels)
             mask = indexer == -1
             if mask.any():
@@ -1532,6 +1532,7 @@ def drop(self, labels, level=None):
         return self.delete(inds)
 
     def _drop_from_level(self, labels, level):
+        labels = com._index_labels_to_array(labels)
         i = self._get_level_number(level)
         index = self.levels[i]
         values = index.get_indexer(labels)

pandas/tests/test_index.py

@@ -395,6 +395,10 @@ def test_drop(self):
 
         self.assertRaises(ValueError, self.strIndex.drop, ['foo', 'bar'])
 
+        dropped = self.strIndex.drop(self.strIndex[0])
+        expected = self.strIndex[1:]
+        self.assert_(dropped.equals(expected))
+
     def test_tuple_union_bug(self):
         import pandas
         import numpy as np

pandas/tests/test_multilevel.py

@@ -1192,6 +1192,14 @@ def test_mixed_depth_drop(self):
         expected = expected.drop([('top','OD','wy')], axis=1)
         assert_frame_equal(expected, result)
 
+        result = df.drop('top', axis=1)
+        assert_frame_equal(expected, result)
+
+        result = df.drop('result1', level=1, axis=1)
+        expected = df.drop([('routine1', 'result1', ''),
+                            ('routine2', 'result1', '')], axis=1)
+        assert_frame_equal(expected, result)
+
     def test_mixed_depth_pop(self):
         arrays = [[  'a', 'top', 'top', 'routine1', 'routine1', 'routine2'],
                   [   '',  'OD',  'OD', 'result1',   'result2',  'result1'],