API/ERR: allow iterators in df.set_index & improve errors

pandas/core/frame.py

@@ -4135,13 +4135,8 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
                    'array, or a list containing only valid column keys and '
                    'one-dimensional arrays.')
 
-        if (is_scalar(keys) or isinstance(keys, tuple)
-                or isinstance(keys, (ABCIndexClass, ABCSeries, np.ndarray))):
-            # make sure we have a container of keys/arrays we can iterate over
-            # tuples can appear as valid column keys!
+        if not isinstance(keys, list):
             keys = [keys]
-        elif not isinstance(keys, list):
-            raise ValueError(err_msg)
 
         missing = []
         for col in keys:
@@ -4150,10 +4145,20 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
                 # tuples are always considered keys, never as list-likes
                 if col not in self:
                     missing.append(col)
-            elif (not isinstance(col, (ABCIndexClass, ABCSeries,
-                                       np.ndarray, list))
-                  or getattr(col, 'ndim', 1) > 1):
-                raise ValueError(err_msg)
+            elif isinstance(col, (ABCIndexClass, ABCSeries,
+                                  np.ndarray, list)):
+                # arrays are fine as long as they are one-dimensional
+                if getattr(col, 'ndim', 1) > 1:
+                    raise ValueError(err_msg)
+            else:
+                # everything else gets tried as a key; see GH 24969
+                try:
+                    self[col]
+                    str(col)
+                except KeyError:
+                    tipo = type(col)
+                    raise ValueError(err_msg,
+                                     'Received column of type {}'.format(tipo))
 
         if missing:
             raise KeyError('{}'.format(missing))

pandas/tests/frame/test_alter_axes.py

@@ -270,6 +270,31 @@ def test_set_index_raise_on_type(self, frame_of_index_cols, box,
             df.set_index(['A', df['A'], box(df['A'])],
                          drop=drop, append=append)
 
+    def test_set_index_custom_label_type(self):
+        # GH 24969
+
+        class Thing:
+            def __init__(self, name, color):
+                self.name = name
+                self.color = color
+
+            def __str__(self):
+                return "<Thing %r>" % (self.name,)
+
+        thing1 = Thing('One', 'red')
+        thing2 = Thing('Two', 'blue')
+        df = DataFrame({thing1: [0, 1], thing2: [2, 3]})
+        expected = DataFrame({thing1: [0, 1]},
+                             index=Index([2, 3], name=thing2))
+
+        # use custom label directly
+        result = df.set_index(thing2)
+        tm.assert_frame_equal(result, expected)
+
+        # custom label wrapped in list
+        result = df.set_index([thing2])
+        tm.assert_frame_equal(result, expected)
+
     def test_construction_with_categorical_index(self):
         ci = tm.makeCategoricalIndex(10)
         ci.name = 'B'