TST/CLN: break up & parametrize tests for df.set_index

Author: h-vetinari

pandas/core/frame.py

@@ -3862,10 +3862,29 @@ def set_index(self, keys, drop=True, append=False, inplace=False,
         -------
         dataframe : DataFrame
         """
-        inplace = validate_bool_kwarg(inplace, 'inplace')
+        from pandas import Series
+
         if not isinstance(keys, list):
             keys = [keys]
 
+        # collect elements from "keys" that are not allowed array types
+        col_labels = [x for x in keys
+                      if not isinstance(x, (Series, Index, MultiIndex,
+                                            list, np.ndarray))]
+        if any(x not in self for x in col_labels):
+            # if there are any labels that are invalid, we raise a KeyError
+            missing = [x for x in col_labels if x not in self]
+            raise KeyError('{}'.format(missing))
+
+        elif len(set(col_labels)) < len(col_labels):
+            # if all are valid labels, but there are duplicates
+            dup = Series(col_labels)
+            dup = list(dup.loc[dup.duplicated()])
+            raise ValueError('Passed duplicate column names '
+                             'to keys: {dup}'.format(dup=dup))
+
+        inplace = validate_bool_kwarg(inplace, 'inplace')
+
         if inplace:
             frame = self
         else:

pandas/tests/frame/conftest.py

@@ -0,0 +1,121 @@
+import pytest
+
+import numpy as np
+
+from pandas import compat
+import pandas.util.testing as tm
+from pandas import DataFrame, date_range, NaT
+
+
+@pytest.fixture
+def frame():
+    return DataFrame(tm.getSeriesData())
+
+
+@pytest.fixture
+def frame2():
+    return DataFrame(tm.getSeriesData(), columns=['D', 'C', 'B', 'A'])
+
+
+@pytest.fixture
+def intframe():
+    df = DataFrame({k: v.astype(int)
+                   for k, v in compat.iteritems(tm.getSeriesData())})
+    # force these all to int64 to avoid platform testing issues
+    return DataFrame({c: s for c, s in compat.iteritems(df)}, dtype=np.int64)
+
+
+@pytest.fixture
+def tsframe():
+    return DataFrame(tm.getTimeSeriesData())
+
+
+@pytest.fixture
+def mixed_frame():
+    df = DataFrame(tm.getSeriesData())
+    df['foo'] = 'bar'
+    return df
+
+
+@pytest.fixture
+def mixed_float():
+    df = DataFrame(tm.getSeriesData())
+    df.A = df.A.astype('float16')
+    df.B = df.B.astype('float32')
+    df.C = df.C.astype('float64')
+    return df
+
+
+@pytest.fixture
+def mixed_float2():
+    df = DataFrame(tm.getSeriesData())
+    df.D = df.D.astype('float16')
+    df.C = df.C.astype('float32')
+    df.B = df.B.astype('float64')
+    return df
+
+
+@pytest.fixture
+def mixed_int():
+    df = DataFrame({k: v.astype(int)
+                   for k, v in compat.iteritems(tm.getSeriesData())})
+    df.A = df.A.astype('uint8')
+    df.B = df.B.astype('int32')
+    df.C = df.C.astype('int64')
+    df.D = np.ones(len(df.D), dtype='uint64')
+    return df
+
+
+@pytest.fixture
+def all_mixed():
+    return DataFrame({'a': 1., 'b': 2, 'c': 'foo',
+                      'float32': np.array([1.] * 10, dtype='float32'),
+                      'int32': np.array([1] * 10, dtype='int32')},
+                     index=np.arange(10))
+
+
+@pytest.fixture
+def tzframe():
+    df = DataFrame({'A': date_range('20130101', periods=3),
+                    'B': date_range('20130101', periods=3,
+                                    tz='US/Eastern'),
+                    'C': date_range('20130101', periods=3,
+                                    tz='CET')})
+    df.iloc[1, 1] = NaT
+    df.iloc[1, 2] = NaT
+    return df
+
+
+@pytest.fixture
+def empty():
+    return DataFrame({})
+
+
+@pytest.fixture
+def ts1():
+    return tm.makeTimeSeries(nper=30)
+
+
+@pytest.fixture
+def ts2():
+    return tm.makeTimeSeries(nper=30)[5:]
+
+
+@pytest.fixture
+def simple():
+    arr = np.array([[1., 2., 3.],
+                    [4., 5., 6.],
+                    [7., 8., 9.]])
+
+    return DataFrame(arr, columns=['one', 'two', 'three'],
+                     index=['a', 'b', 'c'])
+
+
+@pytest.fixture
+def frame_of_index_cols():
+    df = DataFrame({'A': ['foo', 'foo', 'foo', 'bar', 'bar'],
+                    'B': ['one', 'two', 'three', 'one', 'two'],
+                    'C': ['a', 'b', 'c', 'd', 'e'],
+                    'D': np.random.randn(5),
+                    'E': np.random.randn(5)})
+    return df