common.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. """
  2. Boilerplate functions used in defining binary operations.
  3. """
  4. from __future__ import annotations
  5. from functools import wraps
  6. import sys
  7. from typing import Callable
  8. from pandas._libs.lib import item_from_zerodim
  9. from pandas._libs.missing import is_matching_na
  10. from pandas._typing import F
  11. from pandas.core.dtypes.generic import (
  12. ABCDataFrame,
  13. ABCIndex,
  14. ABCSeries,
  15. )
  16. def unpack_zerodim_and_defer(name: str) -> Callable[[F], F]:
  17. """
  18. Boilerplate for pandas conventions in arithmetic and comparison methods.
  19. Parameters
  20. ----------
  21. name : str
  22. Returns
  23. -------
  24. decorator
  25. """
  26. def wrapper(method: F) -> F:
  27. return _unpack_zerodim_and_defer(method, name)
  28. return wrapper
  29. def _unpack_zerodim_and_defer(method, name: str):
  30. """
  31. Boilerplate for pandas conventions in arithmetic and comparison methods.
  32. Ensure method returns NotImplemented when operating against "senior"
  33. classes. Ensure zero-dimensional ndarrays are always unpacked.
  34. Parameters
  35. ----------
  36. method : binary method
  37. name : str
  38. Returns
  39. -------
  40. method
  41. """
  42. if sys.version_info < (3, 9):
  43. from pandas.util._str_methods import (
  44. removeprefix,
  45. removesuffix,
  46. )
  47. stripped_name = removesuffix(removeprefix(name, "__"), "__")
  48. else:
  49. stripped_name = name.removeprefix("__").removesuffix("__")
  50. is_cmp = stripped_name in {"eq", "ne", "lt", "le", "gt", "ge"}
  51. @wraps(method)
  52. def new_method(self, other):
  53. if is_cmp and isinstance(self, ABCIndex) and isinstance(other, ABCSeries):
  54. # For comparison ops, Index does *not* defer to Series
  55. pass
  56. else:
  57. for cls in [ABCDataFrame, ABCSeries, ABCIndex]:
  58. if isinstance(self, cls):
  59. break
  60. if isinstance(other, cls):
  61. return NotImplemented
  62. other = item_from_zerodim(other)
  63. return method(self, other)
  64. return new_method
  65. def get_op_result_name(left, right):
  66. """
  67. Find the appropriate name to pin to an operation result. This result
  68. should always be either an Index or a Series.
  69. Parameters
  70. ----------
  71. left : {Series, Index}
  72. right : object
  73. Returns
  74. -------
  75. name : object
  76. Usually a string
  77. """
  78. if isinstance(right, (ABCSeries, ABCIndex)):
  79. name = _maybe_match_name(left, right)
  80. else:
  81. name = left.name
  82. return name
  83. def _maybe_match_name(a, b):
  84. """
  85. Try to find a name to attach to the result of an operation between
  86. a and b. If only one of these has a `name` attribute, return that
  87. name. Otherwise return a consensus name if they match or None if
  88. they have different names.
  89. Parameters
  90. ----------
  91. a : object
  92. b : object
  93. Returns
  94. -------
  95. name : str or None
  96. See Also
  97. --------
  98. pandas.core.common.consensus_name_attr
  99. """
  100. a_has = hasattr(a, "name")
  101. b_has = hasattr(b, "name")
  102. if a_has and b_has:
  103. try:
  104. if a.name == b.name:
  105. return a.name
  106. elif is_matching_na(a.name, b.name):
  107. # e.g. both are np.nan
  108. return a.name
  109. else:
  110. return None
  111. except TypeError:
  112. # pd.NA
  113. if is_matching_na(a.name, b.name):
  114. return a.name
  115. return None
  116. except ValueError:
  117. # e.g. np.int64(1) vs (np.int64(1), np.int64(2))
  118. return None
  119. elif a_has:
  120. return a.name
  121. elif b_has:
  122. return b.name
  123. return None