Skip to content

Commit

Permalink
cross type arithmetic
Browse files Browse the repository at this point in the history
  • Loading branch information
jreback committed May 23, 2018
1 parent 9b0ccea commit c98667c
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 2 deletions.
9 changes: 7 additions & 2 deletions pandas/core/arrays/integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ def coerce_to_array(values, dtype, mask=None, copy=False):
'mixed-integer', 'mixed-integer-float']:
raise TypeError("{} cannot be converted to an IntegerDtype".format(
values.dtype))

elif not (is_integer_dtype(values) or is_float_dtype(values)):
raise TypeError("{} cannot be converted to an IntegerDtype".format(
values.dtype))
Expand All @@ -126,15 +127,19 @@ def coerce_to_array(values, dtype, mask=None, copy=False):
if not mask.ndim == 1:
raise TypeError("mask must be a 1D list-like")

# avoid float->int numpy conversion issues
if is_object_dtype(values):
mask |= isna(values)

# we copy as need to coerce here
if mask.any():
# we copy as need to coerce here
values = values.copy()
values[mask] = 1

values = values.astype(dtype.type)

else:
values = values.astype(dtype.type, copy=False)

return values, mask


Expand Down
19 changes: 19 additions & 0 deletions pandas/tests/extension/integer/test_integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -393,3 +393,22 @@ def test_to_integer_array(values, expected):
# convert existing arrays to IntegerArrays
result = to_integer_array(values)
tm.assert_extension_array_equal(result, expected)


def test_cross_type_arithmetic():

df = pd.DataFrame({'A': pd.Series([1, 2, np.nan], dtype='Int64'),
'B': pd.Series([1, np.nan, 3], dtype='UInt8'),
'C': [1, 2, 3]})

result = df.A + df.C
expected = pd.Series([2, 4, np.nan], dtype='Int64')
tm.assert_series_equal(result, expected)

result = (df.A + df.C) * 3 == 12
expected = pd.Series([False, True, False])
tm.assert_series_equal(result, expected)

result = df.A + df.B
expected = pd.Series([2, np.nan, np.nan], dtype='Int64')
tm.assert_series_equal(result, expected)

0 comments on commit c98667c

Please sign in to comment.