|
| 1 | +from itertools import product |
1 | 2 | import operator
|
2 | 3 | import re
|
3 | 4 |
|
@@ -66,7 +67,7 @@ def run_arithmetic(self, df, other, assert_func, check_dtype=False, test_flex=Tr
|
66 | 67 | operator_name = "truediv"
|
67 | 68 |
|
68 | 69 | if test_flex:
|
69 |
| - op = lambda x, y: getattr(df, arith)(y) |
| 70 | + op = lambda x, y: getattr(x, arith)(y) |
70 | 71 | op.__name__ = arith
|
71 | 72 | else:
|
72 | 73 | op = getattr(operator, operator_name)
|
@@ -431,3 +432,30 @@ def test_bool_ops_column_name_dtype(self, test_input, expected):
|
431 | 432 | # GH 22383 - .ne fails if columns containing column name 'dtype'
|
432 | 433 | result = test_input.loc[:, ["a", "dtype"]].ne(test_input.loc[:, ["a", "dtype"]])
|
433 | 434 | assert_frame_equal(result, expected)
|
| 435 | + |
| 436 | + @pytest.mark.parametrize( |
| 437 | + "axis,arith", product((0, 1), ("add", "sub", "mul", "mod", "truediv")) |
| 438 | + ) |
| 439 | + def test_frame_series_axis(self, axis, arith): |
| 440 | + # Can't check floordiv here because it currently doesn't work #GH27636 |
| 441 | + df = self.frame |
| 442 | + if axis == 1: |
| 443 | + other = self.frame.iloc[0, :] |
| 444 | + else: |
| 445 | + other = self.frame.iloc[:, 0] |
| 446 | + |
| 447 | + expr._MIN_ELEMENTS = 0 |
| 448 | + |
| 449 | + op = lambda x, y: getattr(x, arith)(y, axis=axis) |
| 450 | + op.__name__ = arith |
| 451 | + |
| 452 | + expr.set_use_numexpr(False) |
| 453 | + expected = op(df, other) |
| 454 | + expr.set_use_numexpr(True) |
| 455 | + |
| 456 | + result = op(df, other) |
| 457 | + try: |
| 458 | + assert_frame_equal(expected, result) |
| 459 | + except Exception: |
| 460 | + pprint_thing("Failed test with operator {op.__name__!r}".format(op=op)) |
| 461 | + raise |
0 commit comments