fixed DataFrame.toCSV bug, test coverage

Author: wesm

pandas/core/frame.py

@@ -738,13 +738,18 @@ def toCSV(self, path, nanRep='', cols=None, header=True,
         for idx in self.index:
             if index:
                 f.write(str(idx))
-            for col in cols:
+            for i, col in enumerate(cols):
                 val = series[col].get(idx)
                 if isnull(val):
                     val = nanRep
                 else:
                     val = str(val)
-                f.write(',%s' % val)
+
+                if i > 0 or index:
+                    f.write(',%s' % val)
+                else:
+                    f.write('%s' % val)
+
             f.write('\n')
 
         f.close()

pandas/core/groupby.py

@@ -337,44 +337,38 @@ def transform(self, func):
         # DataMatrix objects?
         result_values = np.empty_like(self.obj.values)
 
-        if self.axis == 1:
-            result_values = result_values.T
+        if self.axis == 0:
+            trans = lambda x: x
+        elif self.axis == 1:
+            trans = lambda x: x.T
+
+        result_values = trans(result_values)
 
-        # result = {}
         for val, group in self.groups.iteritems():
-            if not isinstance(group, list):
+            if not isinstance(group, list): # pragma: no cover
                 group = list(group)
 
-            subframe = self.obj.reindex(group)
-            subframe.groupName = val
-
             if self.axis == 0:
+                subframe = self.obj.reindex(group)
                 indexer, _ = common.get_indexer(self.obj.index,
                                                 subframe.index, None)
             else:
+                subframe = self.obj.reindex(columns=group)
                 indexer, _ = common.get_indexer(self.obj.columns,
                                                 subframe.columns, None)
+            subframe.groupName = val
 
             try:
-                res = func(subframe)
-            except Exception:
                 res = subframe.apply(func, axis=self.axis)
+            except Exception: # pragma: no cover
+                res = func(subframe)
 
-            result_values[indexer] = res.values
-
-            # result[val] = res
+            result_values[indexer] = trans(res.values)
 
-        if self.axis == 1:
-            result_values = result_values.T
+        result_values = trans(result_values)
 
         return DataFrame(result_values, index=self.obj.index,
                          columns=self.obj.columns)
-        # allSeries = {}
-        # for val, frame in result.iteritems():
-        #     allSeries.update(frame._series)
-
-        # return self._klass(data=allSeries).T
-
 
 class DataMatrixGroupBy(DataFrameGroupBy):
     _klass = DataMatrix

pandas/core/matrix.py

@@ -791,7 +791,7 @@ def info(self, buffer=sys.stdout):
 #-------------------------------------------------------------------------------
 # Public methods
 
-    def apply(self, func, axis=0):
+    def apply(self, func, axis=0, broadcast=False):
         """
         Applies func to columns (Series) of this DataMatrix and returns either
         a DataMatrix (if the function produces another series) or a Series
@@ -802,6 +802,9 @@ def apply(self, func, axis=0):
         ----------
         func : function
             Function to apply to each column
+        broadcast : bool, default False
+            For aggregation functions, return object of same size with values
+            propagated
 
         Examples
         --------
@@ -819,7 +822,8 @@ def apply(self, func, axis=0):
             return DataMatrix(data=results, index=self.index,
                               columns=self.columns, objects=self.objects)
         else:
-            return DataFrame.apply(self, func, axis=axis)
+            return DataFrame.apply(self, func, axis=axis,
+                                   broadcast=broadcast)
 
     def applymap(self, func):
         """

pandas/core/tests/test_frame.py

@@ -696,6 +696,12 @@ def test_toCSV_fromcsv(self):
         recons = self.klass.fromcsv(path, index_col=None)
         assert(len(recons.cols()) == len(self.tsframe.cols()) + 1)
 
+
+        # no index
+        self.tsframe.toCSV(path, index=False)
+        recons = self.klass.fromcsv(path, index_col=None)
+        assert_almost_equal(self.tsframe.values, recons.values)
+
         os.remove(path)
 
     def test_toDataMatrix(self):
@@ -1173,6 +1179,19 @@ def test_apply(self):
         applied = self.empty.apply(np.mean)
         self.assert_(not applied)
 
+
+    def test_apply_broadcast(self):
+        broadcasted = self.frame.apply(np.mean, broadcast=True)
+        agged = self.frame.apply(np.mean)
+
+        for col, ts in broadcasted.iteritems():
+            self.assert_((ts == agged[col]).all())
+
+        broadcasted = self.frame.apply(np.mean, axis=1, broadcast=True)
+        agged = self.frame.apply(np.mean, axis=1)
+        for idx in broadcasted.index:
+            self.assert_((broadcasted.xs(idx) == agged[idx]).all())
+
     def test_tapply(self):
         d = self.frame.index[0]
         tapplied = self.frame.tapply(np.mean)
@@ -1234,6 +1253,11 @@ def test_groupby_columns(self):
         self.assertEqual(len(aggregated), len(self.tsframe))
         self.assertEqual(len(aggregated.cols()), 2)
 
+        # transform
+        tf = lambda x: x - x.mean()
+        groupedT = self.tsframe.T.groupby(mapping, axis=0)
+        assert_frame_equal(groupedT.transform(tf).T, grouped.transform(tf))
+
         # iterate
         for k, v in grouped:
             self.assertEqual(len(v.cols()), 2)

pandas/core/tests/test_series.py

@@ -825,6 +825,10 @@ def test_groupby(self):
         self.assertEqual(agged[1], 1)
 
         assert_series_equal(agged, grouped.agg(np.mean)) # shorthand
+        assert_series_equal(agged, grouped.mean())
+
+        assert_series_equal(grouped.agg(np.sum), grouped.sum())
+
 
         transformed = grouped.transform(lambda x: x * x.sum())
         self.assertEqual(transformed[7], 12)
@@ -867,5 +871,6 @@ def test_groupby_transform(self):
             for idx in group.index:
                 self.assertEqual(transformed[idx], mean)
 
+
 if __name__ == '__main__':
     unittest.main()