diff --git a/pandas/_libs/lib.pyx b/pandas/_libs/lib.pyx index 0b0334d52c1e9..599667a093290 100644 --- a/pandas/_libs/lib.pyx +++ b/pandas/_libs/lib.pyx @@ -2115,7 +2115,8 @@ def maybe_convert_numeric(ndarray[object] values, set na_values, def maybe_convert_objects(ndarray[object] objects, bint try_float=False, bint safe=False, bint convert_datetime=False, bint convert_timedelta=False, - bint convert_to_nullable_integer=False): + bint convert_to_nullable_integer=False, + bint convert_intenum=False): """ Type inference function-- convert object array to proper dtype @@ -2217,6 +2218,9 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=False, seen.object_ = True break elif util.is_integer_object(val): + if getattr(val, 'name', None) is not None and not convert_intenum: + seen.object_ = True + break seen.int_ = True floats[i] = val complexes[i] = val diff --git a/pandas/_libs/tslibs/util.pxd b/pandas/_libs/tslibs/util.pxd index 16d801f69df05..7cc51fea62afd 100644 --- a/pandas/_libs/tslibs/util.pxd +++ b/pandas/_libs/tslibs/util.pxd @@ -69,7 +69,7 @@ cdef inline bint is_integer_object(object obj) nogil: Notes ----- - This counts np.timedelta64 objects as integers. + This counts np.timedelta64 and IntEnums objects as integers. """ return (not PyBool_Check(obj) and PyArray_IsIntegerScalar(obj) and not is_timedelta64_object(obj)) diff --git a/pandas/tests/dtypes/test_inference.py b/pandas/tests/dtypes/test_inference.py index 438a22c99a4eb..ef5ef8d7a7d1c 100644 --- a/pandas/tests/dtypes/test_inference.py +++ b/pandas/tests/dtypes/test_inference.py @@ -7,6 +7,7 @@ from collections import namedtuple from datetime import date, datetime, time, timedelta from decimal import Decimal +from enum import IntEnum from fractions import Fraction from io import StringIO from numbers import Number @@ -589,6 +590,24 @@ def test_maybe_convert_objects_bool_nan(self): out = lib.maybe_convert_objects(ind.values, safe=1) tm.assert_numpy_array_equal(out, exp) + def test_maybe_convert_objects_intenum(self): + class Colors(IntEnum): + red = 1 + blue = 2 + + ind = Index([Colors.red, Colors.blue], dtype=object) + expected = np.array([Colors.red, Colors.blue], dtype=object) + result = lib.maybe_convert_objects(ind.values) + + # by default, we should not convert IntEnums to ints + tm.assert_numpy_array_equal(result, expected) + + expected = np.array([1, 2], dtype=np.int64) + result = lib.maybe_convert_objects(ind.values, convert_intenum=True) + + # still coverts to int if convert_intenum set to True + tm.assert_numpy_array_equal(result, expected) + def test_mixed_dtypes_remain_object_array(self): # GH14956 array = np.array([datetime(2015, 1, 1, tzinfo=pytz.utc), 1], dtype=object)