Skip to content

Commit

Permalink
Fixturize tests/frame/test_arithmetic
Browse files Browse the repository at this point in the history
  • Loading branch information
@h-vetinari
h-vetinari committed Sep 17, 2018
1 parent 8a1c8ad commit 9ce1860
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 123 deletions.
18 changes: 10 additions & 8 deletions pandas/tests/frame/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,10 @@ def mixed_float_frame():
Columns are ['A', 'B', 'C', 'D'].
"""
df = DataFrame(tm.getSeriesData())
df.A = df.A.astype('float16')
df.A = df.A.astype('float32')
df.B = df.B.astype('float32')
df.C = df.C.astype('float64')
df.C = df.C.astype('float16')
df.D = df.D.astype('float64')
return df


Expand All @@ -84,9 +85,10 @@ def mixed_float_frame2():
Columns are ['A', 'B', 'C', 'D'].
"""
df = DataFrame(tm.getSeriesData())
df.D = df.D.astype('float16')
df.D = df.D.astype('float32')
df.C = df.C.astype('float32')
df.B = df.B.astype('float64')
df.B = df.B.astype('float16')
df.D = df.D.astype('float64')
return df


Expand All @@ -99,10 +101,10 @@ def mixed_int_frame():
"""
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')
df.A = df.A.astype('int32')
df.B = np.ones(len(df.B), dtype='uint64')
df.C = df.C.astype('uint8')
df.D = df.C.astype('int64')
return df


Expand Down
195 changes: 80 additions & 115 deletions pandas/tests/frame/test_arithmetic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,132 +127,97 @@ def test_df_add_flex_filled_mixed_dtypes(self):
'B': ser * 2})
tm.assert_frame_equal(result, expected)

def test_arith_flex_frame(self):
seriesd = tm.getSeriesData()
frame = pd.DataFrame(seriesd).copy()

mixed_float = pd.DataFrame({'A': frame['A'].copy().astype('float32'),
'B': frame['B'].copy().astype('float32'),
'C': frame['C'].copy().astype('float16'),
'D': frame['D'].copy().astype('float64')})

intframe = pd.DataFrame({k: v.astype(int)
for k, v in seriesd.items()})
mixed_int = pd.DataFrame({'A': intframe['A'].copy().astype('int32'),
'B': np.ones(len(intframe), dtype='uint64'),
'C': intframe['C'].copy().astype('uint8'),
'D': intframe['D'].copy().astype('int64')})

# force these all to int64 to avoid platform testing issues
intframe = pd.DataFrame({c: s for c, s in intframe.items()},
dtype=np.int64)

ops = ['add', 'sub', 'mul', 'div', 'truediv', 'pow', 'floordiv', 'mod']
@pytest.mark.parametrize('op', ['add', 'sub', 'mul', 'div', 'truediv',
'pow', 'floordiv', 'mod'])
def test_arith_flex_frame(self, op, int_frame, mixed_int_frame,
float_frame, mixed_float_frame):

if not PY3:
aliases = {}
else:
aliases = {'div': 'truediv'}

for op in ops:
try:
alias = aliases.get(op, op)
f = getattr(operator, alias)
result = getattr(frame, op)(2 * frame)
exp = f(frame, 2 * frame)
alias = aliases.get(op, op)

f = getattr(operator, alias)
result = getattr(float_frame, op)(2 * float_frame)
exp = f(float_frame, 2 * float_frame)
tm.assert_frame_equal(result, exp)

# vs mix float
result = getattr(mixed_float_frame, op)(2 * mixed_float_frame)
exp = f(mixed_float_frame, 2 * mixed_float_frame)
tm.assert_frame_equal(result, exp)
_check_mixed_float(result, dtype=dict(C=None))

# vs mix int
if op in ['add', 'sub', 'mul']:
result = getattr(mixed_int_frame, op)(2 + mixed_int_frame)
exp = f(mixed_int_frame, 2 + mixed_int_frame)

# no overflow in the uint
dtype = None
if op in ['sub']:
dtype = dict(B='uint64', C=None)
elif op in ['add', 'mul']:
dtype = dict(C=None)
tm.assert_frame_equal(result, exp)
_check_mixed_int(result, dtype=dtype)

# rops
r_f = lambda x, y: f(y, x)
result = getattr(float_frame, 'r' + op)(2 * float_frame)
exp = r_f(float_frame, 2 * float_frame)
tm.assert_frame_equal(result, exp)

# vs mix float
result = getattr(mixed_float_frame, op)(2 * mixed_float_frame)
exp = f(mixed_float_frame, 2 * mixed_float_frame)
tm.assert_frame_equal(result, exp)
_check_mixed_float(result, dtype=dict(C=None))

result = getattr(int_frame, op)(2 * int_frame)
exp = f(int_frame, 2 * int_frame)
tm.assert_frame_equal(result, exp)

# vs mix int
if op in ['add', 'sub', 'mul']:
result = getattr(mixed_int_frame, op)(2 + mixed_int_frame)
exp = f(mixed_int_frame, 2 + mixed_int_frame)

# no overflow in the uint
dtype = None
if op in ['sub']:
dtype = dict(B='uint64', C=None)
elif op in ['add', 'mul']:
dtype = dict(C=None)
tm.assert_frame_equal(result, exp)
_check_mixed_int(result, dtype=dtype)

# vs mix float
result = getattr(mixed_float, op)(2 * mixed_float)
exp = f(mixed_float, 2 * mixed_float)
tm.assert_frame_equal(result, exp)
_check_mixed_float(result, dtype=dict(C=None))

# vs mix int
if op in ['add', 'sub', 'mul']:
result = getattr(mixed_int, op)(2 + mixed_int)
exp = f(mixed_int, 2 + mixed_int)

# no overflow in the uint
dtype = None
if op in ['sub']:
dtype = dict(B='uint64', C=None)
elif op in ['add', 'mul']:
dtype = dict(C=None)
tm.assert_frame_equal(result, exp)
_check_mixed_int(result, dtype=dtype)

# rops
r_f = lambda x, y: f(y, x)
result = getattr(frame, 'r' + op)(2 * frame)
exp = r_f(frame, 2 * frame)
tm.assert_frame_equal(result, exp)

# vs mix float
result = getattr(mixed_float, op)(2 * mixed_float)
exp = f(mixed_float, 2 * mixed_float)
tm.assert_frame_equal(result, exp)
_check_mixed_float(result, dtype=dict(C=None))

result = getattr(intframe, op)(2 * intframe)
exp = f(intframe, 2 * intframe)
tm.assert_frame_equal(result, exp)

# vs mix int
if op in ['add', 'sub', 'mul']:
result = getattr(mixed_int, op)(2 + mixed_int)
exp = f(mixed_int, 2 + mixed_int)

# no overflow in the uint
dtype = None
if op in ['sub']:
dtype = dict(B='uint64', C=None)
elif op in ['add', 'mul']:
dtype = dict(C=None)
tm.assert_frame_equal(result, exp)
_check_mixed_int(result, dtype=dtype)
except:
printing.pprint_thing("Failing operation %r" % op)
raise

# ndim >= 3
ndim_5 = np.ones(frame.shape + (3, 4, 5))
msg = "Unable to coerce to Series/DataFrame"
with tm.assert_raises_regex(ValueError, msg):
f(frame, ndim_5)

with tm.assert_raises_regex(ValueError, msg):
getattr(frame, op)(ndim_5)

# res_add = frame.add(frame)
# res_sub = frame.sub(frame)
# res_mul = frame.mul(frame)
# res_div = frame.div(2 * frame)

# tm.assert_frame_equal(res_add, frame + frame)
# tm.assert_frame_equal(res_sub, frame - frame)
# tm.assert_frame_equal(res_mul, frame * frame)
# tm.assert_frame_equal(res_div, frame / (2 * frame))

const_add = frame.add(1)
tm.assert_frame_equal(const_add, frame + 1)
# ndim >= 3
ndim_5 = np.ones(float_frame.shape + (3, 4, 5))
msg = "Unable to coerce to Series/DataFrame"
with tm.assert_raises_regex(ValueError, msg):
f(float_frame, ndim_5)

with tm.assert_raises_regex(ValueError, msg):
getattr(float_frame, op)(ndim_5)

const_add = float_frame.add(1)
tm.assert_frame_equal(const_add, float_frame + 1)

# corner cases
result = frame.add(frame[:0])
tm.assert_frame_equal(result, frame * np.nan)
result = float_frame.add(float_frame[:0])
tm.assert_frame_equal(result, float_frame * np.nan)

result = frame[:0].add(frame)
tm.assert_frame_equal(result, frame * np.nan)
result = float_frame[:0].add(float_frame)
tm.assert_frame_equal(result, float_frame * np.nan)
with tm.assert_raises_regex(NotImplementedError, 'fill_value'):
frame.add(frame.iloc[0], fill_value=3)
float_frame.add(float_frame.iloc[0], fill_value=3)
with tm.assert_raises_regex(NotImplementedError, 'fill_value'):
frame.add(frame.iloc[0], axis='index', fill_value=3)

def test_arith_flex_series(self):
arr = np.array([[1., 2., 3.],
[4., 5., 6.],
[7., 8., 9.]])
df = pd.DataFrame(arr, columns=['one', 'two', 'three'],
index=['a', 'b', 'c'])
float_frame.add(float_frame.iloc[0], axis='index', fill_value=3)

def test_arith_flex_series(self, simple_frame):
df = simple_frame

row = df.xs('a')
col = df['two']
Expand Down

0 comments on commit 9ce1860

Please sign in to comment.