ENH: DataFrame.insert method implementation

Author: wesm

pandas/core/frame.py

@@ -717,7 +717,7 @@ def __setitem__(self, key, value):
 
             self._boolean_set(key, value)
         else:
-            self._insert_item(key, value)
+            self._set_item(key, value)
 
     def _boolean_set(self, key, value):
         mask = key.values
@@ -729,7 +729,23 @@ def _boolean_set(self, key, value):
 
         self.values[mask] = value
 
-    def _insert_item(self, key, value):
+    def insert(self, loc, column, value):
+        """
+        Insert column into DataFrame at specified location. Raises Exception if
+        column is already contained in the DataFrame
+
+        Parameters
+        ----------
+        loc : int
+            Must have 0 <= loc <= len(columns)
+        column : object
+        value : int, Series, or array-like
+        """
+        value = self._sanitize_column(value)
+        value = np.atleast_2d(value) # is this a hack?
+        self._data.insert(loc, column, value)
+
+    def _set_item(self, key, value):
         """
         Add series to DataFrame in specified column.
 
@@ -739,9 +755,13 @@ def _insert_item(self, key, value):
         Series/TimeSeries will be conformed to the DataFrame's index to
         ensure homogeneity.
         """
+        value = self._sanitize_column(value)
+        value = np.atleast_2d(value) # is this a hack?
+        self._data.set(key, value)
+
+    def _sanitize_column(self, value):
         # Need to make sure new columns (which go into the BlockManager as new
         # blocks) are always copied
-
         if hasattr(value, '__iter__'):
             if isinstance(value, Series):
                 if value.index.equals(self.index):
@@ -761,8 +781,7 @@ def _insert_item(self, key, value):
         else:
             value = np.repeat(value, len(self.index))
 
-        value = np.atleast_2d(value)
-        self._data.set(key, value)
+        return value
 
     def __delitem__(self, key):
         """

pandas/core/internals.py

@@ -490,9 +490,9 @@ def set(self, item, value):
                 block.set(item, value)
         else:
             # insert at end
-            self.insert(item, len(self.items), value)
+            self.insert(len(self.items), item, value)
 
-    def insert(self, item, loc, value):
+    def insert(self, loc, item, value):
         if item in self.items:
             raise Exception('cannot insert %s, already exists' % item)
 

pandas/tests/test_frame.py

@@ -977,6 +977,21 @@ def test_toString(self):
         frame = self.klass(index=np.arange(1000))
         frame.toString(buf=buf)
 
+    def test_insert(self):
+        df = DataFrame(np.random.randn(5, 3), index=np.arange(5),
+                       columns=['c', 'b', 'a'])
+
+        df.insert(0, 'foo', df['a'])
+        self.assert_(np.array_equal(df.columns, ['foo', 'c', 'b', 'a']))
+        assert_almost_equal(df['a'], df['foo'])
+
+        df.insert(2, 'bar', df['c'])
+        self.assert_(np.array_equal(df.columns, ['foo', 'c', 'bar', 'b', 'a']))
+        assert_almost_equal(df['c'], df['bar'])
+
+        self.assertRaises(Exception, df.insert, 1, 'a')
+        self.assertRaises(Exception, df.insert, 1, 'c')
+
     def test_delitem(self):
         del self.frame['A']
         self.assert_('A' not in self.frame)