Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Functions for defining unary operations. 

3""" 

4from typing import Any, Union 

5 

6import numpy as np 

7 

8from pandas.core.dtypes.common import ( 

9 is_datetime64_dtype, 

10 is_extension_array_dtype, 

11 is_integer_dtype, 

12 is_object_dtype, 

13 is_scalar, 

14 is_timedelta64_dtype, 

15) 

16from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries 

17 

18from pandas.core.construction import array 

19 

20 

21def should_extension_dispatch(left: ABCSeries, right: Any) -> bool: 

22 """ 

23 Identify cases where Series operation should use dispatch_to_extension_op. 

24 

25 Parameters 

26 ---------- 

27 left : Series 

28 right : object 

29 

30 Returns 

31 ------- 

32 bool 

33 """ 

34 if ( 

35 is_extension_array_dtype(left.dtype) 

36 or is_datetime64_dtype(left.dtype) 

37 or is_timedelta64_dtype(left.dtype) 

38 ): 

39 return True 

40 

41 if not is_scalar(right) and is_extension_array_dtype(right): 

42 # GH#22378 disallow scalar to exclude e.g. "category", "Int64" 

43 return True 

44 

45 return False 

46 

47 

48def should_series_dispatch(left, right, op): 

49 """ 

50 Identify cases where a DataFrame operation should dispatch to its 

51 Series counterpart. 

52 

53 Parameters 

54 ---------- 

55 left : DataFrame 

56 right : DataFrame or Series 

57 op : binary operator 

58 

59 Returns 

60 ------- 

61 override : bool 

62 """ 

63 if left._is_mixed_type or right._is_mixed_type: 

64 return True 

65 

66 if op.__name__.strip("_") in ["and", "or", "xor", "rand", "ror", "rxor"]: 

67 # TODO: GH references for what this fixes 

68 # Note: this check must come before the check for nonempty columns. 

69 return True 

70 

71 if right.ndim == 1: 

72 # operating with Series, short-circuit checks that would fail 

73 # with AttributeError. 

74 return False 

75 

76 if not len(left.columns) or not len(right.columns): 

77 # ensure obj.dtypes[0] exists for each obj 

78 return False 

79 

80 ldtype = left.dtypes.iloc[0] 

81 rdtype = right.dtypes.iloc[0] 

82 

83 if (is_timedelta64_dtype(ldtype) and is_integer_dtype(rdtype)) or ( 

84 is_timedelta64_dtype(rdtype) and is_integer_dtype(ldtype) 

85 ): 

86 # numpy integer dtypes as timedelta64 dtypes in this scenario 

87 return True 

88 

89 if is_datetime64_dtype(ldtype) and is_object_dtype(rdtype): 

90 # in particular case where right is an array of DateOffsets 

91 return True 

92 

93 return False 

94 

95 

96def dispatch_to_extension_op( 

97 op, left: Union[ABCExtensionArray, np.ndarray], right: Any, 

98): 

99 """ 

100 Assume that left or right is a Series backed by an ExtensionArray, 

101 apply the operator defined by op. 

102 

103 Parameters 

104 ---------- 

105 op : binary operator 

106 left : ExtensionArray or np.ndarray 

107 right : object 

108 

109 Returns 

110 ------- 

111 ExtensionArray or np.ndarray 

112 2-tuple of these if op is divmod or rdivmod 

113 """ 

114 # NB: left and right should already be unboxed, so neither should be 

115 # a Series or Index. 

116 

117 if left.dtype.kind in "mM" and isinstance(left, np.ndarray): 

118 # We need to cast datetime64 and timedelta64 ndarrays to 

119 # DatetimeArray/TimedeltaArray. But we avoid wrapping others in 

120 # PandasArray as that behaves poorly with e.g. IntegerArray. 

121 left = array(left) 

122 

123 # The op calls will raise TypeError if the op is not defined 

124 # on the ExtensionArray 

125 res_values = op(left, right) 

126 return res_values