Better error message for DataFrame.apply if axis is not 0 or 1

Author: changhiskhan

pandas/core/frame.py

@@ -3760,6 +3760,8 @@ def _apply_standard(self, func, axis, ignore_failures=False):
             series_gen = (Series.from_array(arr, index=res_columns, name=name)
                           for i, (arr, name) in
                           enumerate(izip(values, res_index)))
+        else:
+            raise ValueError('Axis must be 0 or 1, got %s' % str(axis))
 
         keys = []
         results = {}
@@ -3815,6 +3817,8 @@ def _apply_broadcast(self, func, axis):
             target = self
         elif axis == 1:
             target = self.T
+        else:
+            raise ValueError('Axis must be 0 or 1, got %s' % str(axis))
 
         result_values = np.empty_like(target.values)
         columns = target.columns

pandas/tests/test_frame.py

@@ -5239,6 +5239,10 @@ def test_apply(self):
         expected = Series(np.nan, index=self.frame.index)
         assert_series_equal(result, expected)
 
+        #invalid axis
+        df = DataFrame([[1,2,3], [4,5,6], [7,8,9]], index=['a','a','c'])
+        self.assertRaises(ValueError, df.apply, lambda x: x, 2)
+
     def test_apply_standard_nonunique(self):
         df = DataFrame([[1,2,3], [4,5,6], [7,8,9]], index=['a','a','c'])
         rs = df.apply(lambda s: s[0], axis=1)