updates

Author: TomAugspurger

pandas/core/arrays/boolean.py

@@ -1,5 +1,5 @@
 import numbers
-from typing import TYPE_CHECKING, Type
+from typing import TYPE_CHECKING, Optional, Type, Union
 import warnings
 
 import numpy as np
@@ -565,6 +565,7 @@ def logical_method(self, other):
             assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
             other = lib.item_from_zerodim(other)
             other_is_booleanarray = isinstance(other, BooleanArray)
+            other_is_scalar = lib.is_scalar(other)
             mask = None
 
             if other_is_booleanarray:
@@ -577,7 +578,7 @@ def logical_method(self, other):
                     )
                 other, mask = coerce_to_array(other, copy=False)
 
-            if not lib.is_scalar(other) and len(self) != len(other):
+            if not other_is_scalar and len(self) != len(other):
                 raise ValueError("Lengths must match to compare")
 
             if op.__name__ in {"or_", "ror_"}:
@@ -741,13 +742,38 @@ def boolean_arithmetic_method(self, other):
         return set_function_name(boolean_arithmetic_method, name, cls)
 
 
-def kleene_or(left, right, left_mask, right_mask):
+def kleene_or(
+    left: Union[bool, np.nan, np.ndarray],
+    right: Union[bool, np.nan, np.ndarary],
+    left_mask: Optional[np.ndarary],
+    right_mask: Optional[np.ndarray],
+):
+    """
+    Boolean ``or`` using Kleene logic.
+
+    Values are NA where we have ``NA | NA`` or ``NA | False``.
+    ``NA | True`` is considered True.
+
+    Parameters
+    ----------
+    left, right : ndarray, NA, or bool
+        The values of the array.
+    left_mask, right_mask : ndarray, optional
+        The masks. When
+
+    Returns
+    -------
+    result, mask: ndarray[bool]
+        The result of the logical or, and the new mask.
+    """
+    # To reduce the number of cases, we ensure that `left` & `left_mask`
+    # always come from an array, not a scalar. This is safe, since because
+    # A | B == B | A
     if left_mask is None:
         return kleene_or(right, left, right_mask, left_mask)
 
     assert left_mask is not None
-    assert isinstance(left, np.ndarray)
-    assert isinstance(left_mask, np.ndarray)
+    right_is_scalar = right_mask is None
 
     mask = left_mask
 
@@ -757,35 +783,59 @@ def kleene_or(left, right, left_mask, right_mask):
         mask = mask.copy()
 
     # handle scalars:
-    if lib.is_scalar(right) and np.isnan(right):
+    if right_is_scalar and np.isnan(right):  # TODO(pd.NA): change to NA
         result = left.copy()
         mask = left_mask.copy()
         mask[~result] = True
         return result, mask
 
-    # XXX: this implicitly relies on masked values being False!
+    # XXX: verify that this doesn't assume masked values are False!
     result = left | right
     mask[result] = False
 
     # update
     return result, mask
 
 
-def kleene_xor(left, right, left_mask, right_mask):
+def kleene_xor(
+    left: Union[bool, np.nan, np.ndarray],
+    right: Union[bool, np.nan, np.ndarary],
+    left_mask: Optional[np.ndarary],
+    right_mask: Optional[np.ndarray],
+):
+    """
+    Boolean ``xor`` using Kleene logic.
+
+    This is the same as ``or``, with the following adjustments
+
+    * True, True -> False
+    * True, NA   -> NA
+
+    Parameters
+    ----------
+    left, right : ndarray, NA, or bool
+        The values of the array.
+    left_mask, right_mask : ndarray, optional
+        The masks. When
+
+    Returns
+    -------
+    result, mask: ndarray[bool]
+        The result of the logical xor, and the new mask.
+    """
     if left_mask is None:
         return kleene_xor(right, left, right_mask, left_mask)
 
+    # Re-use or, and update with adustments.
     result, mask = kleene_or(left, right, left_mask, right_mask)
-    #
-    # if lib.is_scalar(right):
-    #     if right is True:
-    #         result[result] = False
-    #     result[left & right] = False
 
+    # TODO(pd.NA): change to pd.NA
     if lib.is_scalar(right) and right is np.nan:
+        # True | NA == True
+        # True ^ NA == NA
         mask[result] = True
     else:
-        # assumes masked values are False
+        # XXX: verify that this doesn't assume masked values are False!
         result[left & right] = False
         mask[right & left_mask] = True
         if right_mask is not None:
@@ -795,7 +845,29 @@ def kleene_xor(left, right, left_mask, right_mask):
     return result, mask
 
 
-def kleene_and(left, right, left_mask, right_mask):
+def kleene_and(
+    left: Union[bool, np.nan, np.ndarray],
+    right: Union[bool, np.nan, np.ndarary],
+    left_mask: Optional[np.ndarary],
+    right_mask: Optional[np.ndarray],
+):
+    """
+    Boolean ``and`` using Kleene logic.
+
+    Values are ``NA`` for ``NA & NA`` or ``True & NA``.
+
+    Parameters
+    ----------
+    left, right : ndarray, NA, or bool
+        The values of the array.
+    left_mask, right_mask : ndarray, optional
+        The masks. When
+
+    Returns
+    -------
+    result, mask: ndarray[bool]
+        The result of the logical xor, and the new mask.
+    """
     if left_mask is None:
         return kleene_and(right, left, right_mask, left_mask)