jim22k · Jul 7, 2022 · Jul 4, 2022 · Jul 5, 2022 · Jul 7, 2022
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@
 # To run: `pre-commit run --all-files`
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.2.0
+    rev: v4.3.0
     hooks:
       - id: check-added-large-files
       - id: check-yaml
@@ -23,12 +23,12 @@ repos:
       - id: isort
         language_version: python3
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.32.0
+    rev: v2.34.0
     hooks:
       - id: pyupgrade
         args: [--py38-plus]
   - repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 22.6.0
     hooks:
       - id: black
         language_version: python3

diff --git a/graphblas/base.py b/graphblas/base.py
@@ -282,7 +282,7 @@ def __rand__(self, other):
             return NotImplemented
         from .infix import _ewise_infix_expr
 
-        return _ewise_infix_expr(self, other, method="ewise_mult", within="__rand__")
+        return _ewise_infix_expr(other, self, method="ewise_mult", within="__rand__")
 
     def __matmul__(self, other):
         if self._is_scalar:

diff --git a/graphblas/infix.py b/graphblas/infix.py
@@ -2,6 +2,7 @@
 from .base import _expect_op, _expect_type
 from .dtypes import BOOL
 from .expr import InfixExprBase
+from .mask import Mask
 from .matrix import Matrix, MatrixExpression, TransposedMatrix
 from .monoid import land, lor
 from .scalar import Scalar, ScalarExpression
@@ -352,7 +353,11 @@ def _ewise_infix_expr(left, right, *, method, within):
         elif method == "ewise_mult":
             return MatrixEwiseMultExpr(left, right)
         return MatrixEwiseAddExpr(left, right)
-    if left_type in types:
+    if within == "__or__" and isinstance(right, Mask):
+        return right.__ror__(left)
+    elif within == "__and__" and isinstance(right, Mask):
+        return right.__rand__(left)
+    elif left_type in types:
         left._expect_type(right, tuple(types), within=within, argname="right")
     elif right_type in types:
         right._expect_type(left, tuple(types), within=within, argname="left")

diff --git a/graphblas/mask.py b/graphblas/mask.py
@@ -1,6 +1,6 @@
 import warnings
 
-from . import utils
+from . import monoid, utils
 from .binary import land, lor, pair
 from .dtypes import BOOL
 from .select import valuene
@@ -62,7 +62,7 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None):
 
         >>> val = Matrix(...)
         >>> val(self) << True
-        >>> val(mask, replace=True) << True
+        >>> val(mask, replace=True) << val
 
         If `complement=` argument is True, then the *complement* will be returned.
         This is equivalent to the following (but uses more efficient recipes):
@@ -89,6 +89,54 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None):
         func = d[type(self), type(mask)]
         return func(self, mask, dtype, name)
 
+    def __and__(self, other):
+        """Return the intersection of two masks as a new mask.
+
+        `new_mask = mask1 & mask2` is equivalent to the following:
+
+        >>> val = Matrix(bool, nrows, ncols)
+        >>> val(mask1) << True
+        >>> val(mask2, replace=True) << val
+        >>> new_mask = val.S
+
+        This uses faster recipes than the above for all combinations of input mask types,
+        and aims to be memory efficient when operating on complemented masks.
+        """
+        from .base import _check_mask
+
+        other = _check_mask(other)
+        complement = self.complement or other.complement
+        d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS
+        func = d[type(self), type(other)]
+        val = func(self, other, bool, None)
+        if complement:
+            return ComplementedStructuralMask(val)
+        else:
+            return StructuralMask(val)
+
+    __rand__ = __and__
+
+    def __or__(self, other):
+        """Return the union of two masks as a new mask.
+
+        `new_mask = mask1 | mask2` is equivalent to the following:
+
+        >>> val = Matrix(bool, nrows, ncols)
+        >>> val(mask1) << True
+        >>> val(mask2) << True
+        >>> new_mask = val.S
+
+        This uses faster recipes than the above for all combinations of input mask types,
+        and aims to be memory efficient when operating on complemented masks.
+        """
+        from .base import _check_mask
+
+        other = _check_mask(other)
+        func = _MASK_OR[type(self), type(other)]
+        return func(self, other)
+
+    __ror__ = __or__
+
 
 class StructuralMask(Mask):
     __slots__ = ()
@@ -371,6 +419,97 @@ def _complement_A_CV(m1, m2, dtype, name):
     # CV-CV
     (ComplementedValueMask, ComplementedValueMask): _complement_CV_CV,
 }
+
+
+def _combine_S_S_mask_or(m1, m2):
+    """S-S"""
+    val = monoid.any(one(m1.parent).new(bool) | one(m2.parent).new(bool)).new()
+    return StructuralMask(val)
+
+
+def _combine_S_SV_mask_or(m1, m2):
+    """S-V"""
+    val = monoid.any(one(m1.parent).new(bool) | one(m2.parent).new(bool, mask=m2)).new()
+    return StructuralMask(val)
+
+
+def _combine_SV_S_mask_or(m1, m2):
+    """V-S"""
+    val = monoid.any(one(m1.parent).new(bool, mask=m1) | one(m2.parent).new(bool)).new()
+    return StructuralMask(val)
+
+
+def _complement_A_CS_mask_or(m1, m2):
+    """~S-CS, ~V-CS, ~CV-CS"""
+    val = one(m2.parent).new(bool, mask=~m1)
+    return ComplementedStructuralMask(val)
+
+
+def _complement_CS_A_mask_or(m1, m2):
+    """~CS-S, ~CS-V, ~CS-CV"""
+    val = one(m1.parent).new(bool, mask=~m2)
+    return ComplementedStructuralMask(val)
+
+
+def _complement_A_CV_mask_or(m1, m2):
+    """~S-CV, ~V-CV"""
+    val = valuene(m2.parent).new(bool, mask=~m1)
+    return ComplementedStructuralMask(val)
+
+
+def _complement_CV_A_mask_or(m1, m2):
+    """~CV-S, ~CV-V"""
+    val = valuene(m1.parent).new(bool, mask=~m2)
+    return ComplementedStructuralMask(val)
+
+
+def _combine_V_V_mask_or(m1, m2):
+    """V-V"""
+    val = monoid.any(valuene(m1.parent).new() | valuene(m2.parent).new()).new(bool)
+    return StructuralMask(val)
+
+
+def _complement_CS_CS_mask_or(m1, m2):
+    """~CS-CS"""
+    val = pair(m1.parent & m2.parent).new(bool)
+    return ComplementedStructuralMask(val)
+
+
+def _complement_CV_CV_mask_or(m1, m2):
+    """~CV-CV"""
+    val = valuene(land(m1.parent & m2.parent).new(bool)).new()
+    return ComplementedStructuralMask(val)
+
+
+_MASK_OR = {
+    # S-S
+    (StructuralMask, StructuralMask): _combine_S_S_mask_or,
+    # S-V, V-S
+    (StructuralMask, ValueMask): _combine_S_SV_mask_or,
+    (ValueMask, StructuralMask): _combine_SV_S_mask_or,
+    # S-CS, CS-S
+    (StructuralMask, ComplementedStructuralMask): _complement_A_CS_mask_or,
+    (ComplementedStructuralMask, StructuralMask): _complement_CS_A_mask_or,
+    # S-CV, CV-S
+    (StructuralMask, ComplementedValueMask): _complement_A_CV_mask_or,
+    (ComplementedValueMask, StructuralMask): _complement_CV_A_mask_or,
+    # V-V
+    (ValueMask, ValueMask): _combine_V_V_mask_or,
+    # V-CS, CS-V
+    (ValueMask, ComplementedStructuralMask): _complement_A_CS_mask_or,
+    (ComplementedStructuralMask, ValueMask): _complement_CS_A_mask_or,
+    # V-CV, CV-V
+    (ValueMask, ComplementedValueMask): _complement_A_CV_mask_or,
+    (ComplementedValueMask, ValueMask): _complement_CV_A_mask_or,
+    # CS-CS
+    (ComplementedStructuralMask, ComplementedStructuralMask): _complement_CS_CS_mask_or,
+    # CS-CV, CV-CS
+    (ComplementedStructuralMask, ComplementedValueMask): _complement_CS_A_mask_or,
+    (ComplementedValueMask, ComplementedStructuralMask): _complement_A_CS_mask_or,
+    # CV-CV
+    (ComplementedValueMask, ComplementedValueMask): _complement_CV_CV_mask_or,
+}
+
 utils._output_types[StructuralMask] = StructuralMask
 utils._output_types[ValueMask] = ValueMask
 utils._output_types[ComplementedStructuralMask] = ComplementedStructuralMask

diff --git a/graphblas/tests/test_mask.py b/graphblas/tests/test_mask.py
@@ -3,6 +3,7 @@
 import pytest
 
 from graphblas import Vector
+from graphblas.mask import Mask
 
 
 @pytest.mark.parametrize("as_matrix", [False, True])
@@ -53,3 +54,73 @@ def test_mask_new(as_matrix):
         else:
             with pytest.raises(TypeError, match="Mask must be"):
                 m.new(mask=v1)
+
+
+@pytest.mark.parametrize("as_matrix", [False, True])
+def test_mask_or(as_matrix):
+    for mask_dtype in [bool, int]:
+        v1 = Vector(mask_dtype, size=10)
+        v1[3:6] = 0
+        v1[:3] = 10
+        v2 = Vector(mask_dtype, size=10)
+        v2[1::3] = 0
+        v2[::3] = 10
+        if as_matrix:
+            v1 = v1._as_matrix()
+            v2 = v2._as_matrix()
+        masks = [v1.S, v1.V, ~v1.S, ~v1.V, v2.S, v2.V, ~v2.S, ~v2.V]
+        for m1, m2 in itertools.product(masks, masks):
+            expected = Vector(bool, size=10)
+            if as_matrix:
+                expected = expected._as_matrix()
+            expected(m1) << True
+            expected(m2) << True
+            result = (m1 | m2).new()
+            assert result.isequal(expected, check_dtype=True)
+        with pytest.raises(TypeError, match="Invalid mask"):
+            m1 | object()
+        with pytest.raises(TypeError, match="Invalid mask"):
+            object() | m1
+        if v1.dtype == bool:
+            assert isinstance(m1 | v1, Mask)
+            assert isinstance(v1 | m1, Mask)
+        else:
+            with pytest.raises(TypeError, match="Mask must be"):
+                m1 | v1
+            with pytest.raises(TypeError, match="Mask must be"):
+                v1 | m1
+
+
+@pytest.mark.parametrize("as_matrix", [False, True])
+def test_mask_and(as_matrix):
+    for mask_dtype in [bool, int]:
+        v1 = Vector(mask_dtype, size=10)
+        v1[3:6] = 0
+        v1[:3] = 10
+        v2 = Vector(mask_dtype, size=10)
+        v2[1::3] = 0
+        v2[::3] = 10
+        if as_matrix:
+            v1 = v1._as_matrix()
+            v2 = v2._as_matrix()
+        masks = [v1.S, v1.V, ~v1.S, ~v1.V, v2.S, v2.V, ~v2.S, ~v2.V]
+        for m1, m2 in itertools.product(masks, masks):
+            expected = Vector(bool, size=10)
+            if as_matrix:
+                expected = expected._as_matrix()
+            expected[...] << True
+            expected = expected.dup(mask=m1).dup(mask=m2)
+            result = (m1 & m2).new()
+            assert result.isequal(expected, check_dtype=True)
+        with pytest.raises(TypeError, match="Invalid mask"):
+            m1 & object()
+        with pytest.raises(TypeError, match="Invalid mask"):
+            object() & m1
+        if v1.dtype == bool:
+            assert isinstance(m1 & v1, Mask)
+            assert isinstance(v1 & m1, Mask)
+        else:
+            with pytest.raises(TypeError, match="Mask must be"):
+                m1 & v1
+            with pytest.raises(TypeError, match="Mask must be"):
+                v1 & m1
diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py
@@ -2276,15 +2276,18 @@ def test_import_export_auto(A, do_iso, methods):
             else:
                 d = A2.ss.unpack(format, sort=sort, raw=raw)
             if in_method == "import":
-                import_func = getattr(Matrix.ss, f"import_{import_name}")
+
+                def import_func(x, import_name, **kwargs):
+                    return getattr(Matrix.ss, f"import_{import_name}")(**kwargs)
+
             else:
 
-                def import_func(**kwargs):
-                    getattr(A2.ss, f"pack_{import_name}")(**kwargs)
-                    return A2
+                def import_func(x, import_name, **kwargs):
+                    getattr(x.ss, f"pack_{import_name}")(**kwargs)
+                    return x
 
             d["format"] = import_format
-            other = import_func(take_ownership=take_ownership, **d)
+            other = import_func(A2, import_name, take_ownership=take_ownership, **d)
             if format == "bitmapc" and raw and import_format is None and import_name == "any":
                 # It's 1d, so we can't tell we're column-oriented w/o format keyword
                 assert other.isequal(A_orig.T)
@@ -2293,7 +2296,7 @@ def import_func(**kwargs):
             assert other.ss.is_iso is do_iso
             d["format"] = "bad_format"
             with pytest.raises(ValueError, match="Invalid format"):
-                import_func(**d)
+                import_func(A2, import_name, **d)
     assert A.isequal(A_orig)
     assert A.ss.is_iso is do_iso
     assert A_orig.ss.is_iso is do_iso
@@ -2317,15 +2320,18 @@ def import_func(**kwargs):
             else:
                 d = C2.ss.unpack(format, raw=raw)
             if in_method == "import":
-                import_func = getattr(Matrix.ss, f"import_{import_name}")
+
+                def import_func(x, import_name, **kwargs):
+                    return getattr(Matrix.ss, f"import_{import_name}")(**kwargs)
+
             else:
 
-                def import_func(**kwargs):
-                    getattr(C2.ss, f"pack_{import_name}")(**kwargs)
-                    return C2
+                def import_func(x, import_name, **kwargs):
+                    getattr(x.ss, f"pack_{import_name}")(**kwargs)
+                    return x
 
             d["format"] = import_format
-            other = import_func(take_ownership=take_ownership, **d)
+            other = import_func(C2, import_name, take_ownership=take_ownership, **d)
             if format == "fullc" and raw and import_format is None and import_name == "any":
                 # It's 1d, so we can't tell we're column-oriented w/o format keyword
                 if do_iso:
@@ -2340,7 +2346,7 @@ def import_func(**kwargs):
             assert other.ss.is_iso is do_iso
             d["format"] = "bad_format"
             with pytest.raises(ValueError, match="Invalid format"):
-                import_func(**d)
+                import_func(C2, import_name, **d)
     assert C.isequal(C_orig)
     assert C.ss.is_iso is do_iso
     assert C_orig.ss.is_iso is do_iso