Repository: pyccel/pyccel
Pull request #2358 (2025-06-19 19:54:31)

Add support for bitwise operators with NumPy arrays
Add support for bitwise operators with NumPy arrays. This improves coverage on sections that were introduced for this purpose but were never fully implemented and tested. This is necessary for the tests to pass on #2352

pyccel-bot[bot]: Hello again! Thank you for this new pull request :star_struck:.

Please begin by requesting your checklist using the command `/bot checklist`

EmilyBourne: /bot run pylint pyccel_lint docs spelling linux
EmilyBourne: /bot run pyccel_lint pylint linux
EmilyBourne: /bot run linux
pyccel-bot[bot]: Unfortunately your PR is not passing the tests so it is not quite ready for review yet. Let me know when it is fixed with `/bot mark as ready`.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6168547b..b32f0ac2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -88,8 +88,9 @@ All notable changes to this project will be documented in this file.
 -   #2132 : Add support for `typing.TypeVar` to replace `@template`.
 -   #2253 : Add multiple levels of verbosity.
 -   Generate stub files to allow double compilation to potentially be bypassed.
--   Added `context_dict` argument to `epyccel` for passing non-global `typing.TypeVar` objects.
+-   Add `context_dict` argument to `epyccel` for passing non-global `typing.TypeVar` objects.
 -   #2293 : Add `pyccel-test` command to run unit tests. Improve docs.
+-   #2358 : Add support for bitwise operators with NumPy arrays.
 -   \[INTERNALS\] Add abstract class `SetMethod` to handle calls to various set methods.
 -   \[INTERNALS\] Added `container_rank` property to `ast.datatypes.PyccelType` objects.
 -   \[INTERNALS\] Add a `__call__` method to `FunctionDef` to create `FunctionCall` instances.
diff --git a/pyccel/ast/bitwise_operators.py b/pyccel/ast/bitwise_operators.py
index dbf771df..4dfe5963 100644
--- a/pyccel/ast/bitwise_operators.py
+++ b/pyccel/ast/bitwise_operators.py
@@ -8,10 +8,16 @@ These operators all have a precision as detailed here:
 https://docs.python.org/3/reference/expressions.html#operator-precedence
 They also have specific rules to determine the datatype, rank, shape
 """
-from .builtins     import PythonInt
+from packaging.version import Version
+import numpy
+
+from .builtins     import DtypePrecisionToCastFunction
 from .datatypes    import PrimitiveBooleanType, PrimitiveIntegerType
 from .datatypes    import PythonNativeBool, PythonNativeInt
-from .operators    import PyccelUnaryOperator, PyccelOperator
+from .operators    import PyccelUnaryOperator, PyccelBinaryOperator
+from .numpytypes   import NumpyInt8Type
+
+numpy_v1 = Version(numpy.__version__) < Version("2.0.0")
 
 __all__ = (
     'PyccelBitComparisonOperator',
@@ -63,18 +69,22 @@ class PyccelInvert(PyccelUnaryOperator):
         DataType
             The  datatype of the result of the operation.
         """
-        dtype = PythonNativeInt()
+        if arg.class_type.rank:
+            class_type = arg.class_type
+        else:
+            class_type = arg.class_type.switch_basic_type(PythonNativeInt())
         assert isinstance(getattr(arg.dtype, 'primitive_type', None), (PrimitiveBooleanType, PrimitiveIntegerType))
 
-        self._args      = (PythonInt(arg) if arg.dtype is PythonNativeBool() else arg,)
-        return dtype
+        cast = DtypePrecisionToCastFunction[getattr(class_type, 'element_type', class_type)]
+        self._args = (cast(arg) if arg.dtype is not class_type else arg,)
+        return class_type
 
     def __repr__(self):
         return f'~{repr(self.args[0])}'
 
 #==============================================================================
 
-class PyccelBitOperator(PyccelOperator):
+class PyccelBitOperator(PyccelBinaryOperator):
     """
     Abstract superclass representing a Python bitwise operator with two arguments.
 
@@ -87,8 +97,7 @@ class PyccelBitOperator(PyccelOperator):
     arg2 : TypedAstNode
         The second argument passed to the operator.
     """
-    _shape = None
-    __slots__ = ('_class_type',)
+    __slots__ = ()
 
     def __init__(self, arg1, arg2):
         super().__init__(arg1, arg2)
@@ -119,25 +128,16 @@ class PyccelBitOperator(PyccelOperator):
         DataType
             The  datatype of the result of the operation.
         """
-        try:
-            class_type = arg1.class_type + arg2.class_type
-        except NotImplementedError as err:
-            raise TypeError(f'Cannot determine the type of {arg1} {self.op} {arg2}') from err # pylint: disable=no-member
-
-        assert isinstance(getattr(class_type, 'primitive_type', None), (PrimitiveBooleanType, PrimitiveIntegerType))
+        class_type = arg1.class_type + arg2.class_type
+        if isinstance(getattr(class_type, 'primitive_type', None), PrimitiveBooleanType):
+            assert class_type.rank > 0
+            class_type = class_type.switch_basic_type(NumpyInt8Type())
 
-        self._args = [PythonInt(a) if a.dtype is PythonNativeBool() else a for a in (arg1, arg2)]
+        cast = DtypePrecisionToCastFunction[getattr(class_type, 'element_type', class_type)]
+        self._args = [cast(a) if a.dtype is not class_type else a for a in (arg1, arg2)]
 
         return class_type
 
-    def _set_shape(self):
-        """
-        Set the shape of the resulting object.
-
-        Set the shape of the resulting object. For a PyccelBitOperator,
-        the shape is a class attribute so nothing needs to be done.
-        """
-
     def __repr__(self):
         return f'{self.args[0]} {self.op} {self.args[1]}' # pylint: disable=no-member
 
@@ -231,16 +231,28 @@ class PyccelBitComparisonOperator(PyccelBitOperator):
         DataType
             The  datatype of the result of the operation.
         """
-        try:
-            class_type = arg1.class_type & arg2.class_type
-        except NotImplementedError as err:
-            raise TypeError(f'Cannot determine the type of {arg1} {self.op} {arg2}') from err # pylint: disable=no-member
+        class_type = None
+        if numpy_v1:
+            if arg1.rank > 0 and arg2.rank == 0:
+                class_type = arg1.class_type.switch_basic_type(arg2.class_type) \
+                                if isinstance(arg1.class_type.primitive_type, PrimitiveBooleanType) \
+                                else arg1.class_type
+            elif arg1.rank == 0 and arg2.rank > 0:
+                class_type = arg2.class_type.switch_basic_type(arg1.class_type) \
+                                if isinstance(arg2.class_type.primitive_type, PrimitiveBooleanType) \
+                                else arg2.class_type
+
+        if class_type is None:
+            try:
+                class_type = arg1.class_type & arg2.class_type
+            except NotImplementedError as err:
+                raise TypeError(f'Cannot determine the type of {arg1} {self.op} {arg2}') from err # pylint: disable=no-member
 
         primitive_type = class_type.primitive_type
         assert isinstance(primitive_type, (PrimitiveBooleanType, PrimitiveIntegerType))
+        cast = DtypePrecisionToCastFunction[getattr(class_type, 'element_type', class_type)]
 
-        if isinstance(primitive_type, PrimitiveIntegerType):
-            self._args = [PythonInt(a) if a.dtype is PythonNativeBool() else a for a in (arg1, arg2)]
+        self._args = [cast(a) if a.dtype is not class_type else a for a in (arg1, arg2)]
 
         return class_type
 
diff --git a/pyccel/ast/numpyext.py b/pyccel/ast/numpyext.py
index dc80336a..84f1a3a1 100644
--- a/pyccel/ast/numpyext.py
+++ b/pyccel/ast/numpyext.py
@@ -2759,6 +2759,7 @@ class NumpyNDArray(PyccelFunction):
 #==============================================================================
 
 DtypePrecisionToCastFunction.update({
+    PythonNativeBool()    : NumpyBool,
     NumpyInt8Type()       : NumpyInt8,
     NumpyInt16Type()      : NumpyInt16,
     NumpyInt32Type()      : NumpyInt32,
diff --git a/pyccel/ast/numpytypes.py b/pyccel/ast/numpytypes.py
index e4e7cf18..467014a2 100644
--- a/pyccel/ast/numpytypes.py
+++ b/pyccel/ast/numpytypes.py
@@ -322,9 +322,9 @@ class NumpyNDArrayType(HomogeneousContainerType, metaclass = ArgumentSingleton):
     def __and__(self, other):
         elem_type = self.element_type
         if isinstance(other, FixedSizeNumericType):
-            return NumpyNDArrayType(elem_type and other)
+            return self.switch_basic_type(elem_type & other)
         elif isinstance(other, NumpyNDArrayType):
-            return NumpyNDArrayType(elem_type+other.element_type)
+            return self.switch_basic_type(elem_type & other.element_type)
         else:
             return NotImplemented
 
diff --git a/pyccel/codegen/printing/ccode.py b/pyccel/codegen/printing/ccode.py
index 8076e589..d38383c6 100644
--- a/pyccel/codegen/printing/ccode.py
+++ b/pyccel/codegen/printing/ccode.py
@@ -2516,7 +2516,11 @@ class CCodePrinter(CodePrinter):
         return ' & '.join(args)
 
     def _print_PyccelInvert(self, expr):
-        return '~{}'.format(self._print(expr.args[0]))
+        arg = self._print(expr.args[0])
+        if expr.dtype is PythonNativeBool():
+            return f'!{arg}'
+        else:
+            return f'~{arg}'
 
     def _print_PyccelAssociativeParenthesis(self, expr):
         return '({})'.format(self._print(expr.args[0]))
diff --git a/pyccel/codegen/printing/fcode.py b/pyccel/codegen/printing/fcode.py
index 600d653f..58ca6978 100644
--- a/pyccel/codegen/printing/fcode.py
+++ b/pyccel/codegen/printing/fcode.py
@@ -3086,28 +3086,37 @@ class FCodePrinter(CodePrinter):
         return code
 
     def _print_PyccelRShift(self, expr):
-        return 'RSHIFT({}, {})'.format(self._print(expr.args[0]), self._print(expr.args[1]))
+        arg0, arg1 = expr.args
+        return f'RSHIFT({self._print(arg0)}, {self._print(arg1)})'
 
     def _print_PyccelLShift(self, expr):
-        return 'LSHIFT({}, {})'.format(self._print(expr.args[0]), self._print(expr.args[1]))
+        arg0, arg1 = expr.args
+        return f'LSHIFT({self._print(arg0)}, {self._print(arg1)})'
 
     def _print_PyccelBitXor(self, expr):
         if isinstance(expr.dtype.primitive_type, PrimitiveBooleanType):
             return ' .neqv. '.join(self._print(a) for a in expr.args)
-        return 'IEOR({}, {})'.format(self._print(expr.args[0]), self._print(expr.args[1]))
+        arg0, arg1 = expr.args
+        return f'IEOR({self._print(arg0)}, {self._print(arg1)})'
 
     def _print_PyccelBitOr(self, expr):
         if isinstance(expr.dtype.primitive_type, PrimitiveBooleanType):
             return ' .or. '.join(self._print(a) for a in expr.args)
-        return 'IOR({}, {})'.format(self._print(expr.args[0]), self._print(expr.args[1]))
+        arg0, arg1 = expr.args
+        return f'IOR({self._print(arg0)}, {self._print(arg1)})'
 
     def _print_PyccelBitAnd(self, expr):
         if isinstance(expr.dtype.primitive_type, PrimitiveBooleanType):
             return ' .and. '.join(self._print(a) for a in expr.args)
-        return 'IAND({}, {})'.format(self._print(expr.args[0]), self._print(expr.args[1]))
+        arg0, arg1 = expr.args
+        return f'IAND({self._print(arg0)}, {self._print(arg1)})'
 
     def _print_PyccelInvert(self, expr):
-        return 'NOT({})'.format(self._print(expr.args[0]))
+        arg = self._print(expr.args[0])
+        if expr.dtype is PythonNativeBool():
+            return f'.not. ({arg})'
+        else:
+            return f'NOT({arg})'
 
     def _print_PyccelAssociativeParenthesis(self, expr):
         return '({})'.format(self._print(expr.args[0]))
diff --git a/pyccel/codegen/printing/pycode.py b/pyccel/codegen/printing/pycode.py
index 40c15e4f..2a689bd0 100644
--- a/pyccel/codegen/printing/pycode.py
+++ b/pyccel/codegen/printing/pycode.py
@@ -24,7 +24,7 @@ from pyccel.ast.low_level_tools import UnpackManagedMemory
 from pyccel.ast.numpyext   import numpy_target_swap, numpy_linalg_mod, numpy_random_mod
 from pyccel.ast.numpyext   import NumpyArray, NumpyNonZero, NumpyResultType
 from pyccel.ast.numpyext   import process_dtype as numpy_process_dtype
-from pyccel.ast.numpyext   import NumpyNDArray
+from pyccel.ast.numpyext   import NumpyNDArray, NumpyBool
 from pyccel.ast.numpytypes import NumpyNumericType, NumpyNDArrayType
 from pyccel.ast.type_annotations import VariableTypeAnnotation, SyntacticTypeAnnotation
 from pyccel.ast.typingext  import TypingTypeVar, TypingFinal
@@ -179,6 +179,8 @@ class PythonCodePrinter(CodePrinter):
             cls = expr
         else:
             cls = type(expr)
+        if cls is NumpyBool:
+            return 'bool'
         type_name = expr.name
         name = self._aliases.get(cls, type_name)
         if name == type_name and cls not in (PythonBool, PythonInt, PythonFloat, PythonComplex):


Test patch:
diff --git a/tests/epyccel/recognised_functions/test_numpy_operators.py b/tests/epyccel/recognised_functions/test_numpy_operators.py
new file mode 100644
index 00000000..ec8a202e
--- /dev/null
+++ b/tests/epyccel/recognised_functions/test_numpy_operators.py
@@ -0,0 +1,204 @@
+# pylint: disable=missing-function-docstring, missing-module-docstring
+import os
+from typing import TypeVar
+from numpy.random import randint
+import numpy as np
+import pytest
+
+from pyccel import epyccel
+
+int_types = (bool, np.int8, np.int64)
+IT = TypeVar('IT', *int_types)
+IT2 = TypeVar('IT2', *int_types)
+
+def test_numpy_bit_and(language):
+    def f(a : 'IT[:,:,:]', b : 'IT2[:,:,:]'):
+        return a & b
+    def g(a : 'IT[:,:,:]', b : 'IT2'):
+        return a & b
+    def h(a : 'IT', b : 'IT2[:,:,:]'):
+        return a & b
+
+    epyc_f = epyccel(f, language=language)
+    epyc_g = epyccel(g, language=language)
+    epyc_h = epyccel(h, language=language)
+
+    for t_x in int_types:
+        for t_y in int_types:
+            x = randint(2 if t_x is bool else 127, size=(2,3,4), dtype = t_x)
+            y = randint(2 if t_y is bool else 127, size=(2,3,4), dtype = t_y)
+
+            z = f(x,y)
+            epyc_z = epyc_f(x,y)
+            assert np.array_equal(epyc_z, z)
+            assert z.dtype is epyc_z.dtype
+
+            y = randint(2 if t_y is bool else 127, dtype = t_y)
+
+            z2 = g(x,y)
+            epyc_z2 = epyc_g(x,y)
+            assert np.array_equal(epyc_z2, z2)
+            assert z2.dtype is epyc_z2.dtype
+
+            z3 = h(y,x)
+            epyc_z3 = epyc_h(y,x)
+            assert np.array_equal(epyc_z3, z3)
+            assert z3.dtype is epyc_z3.dtype
+
+def test_numpy_bit_or(language):
+    def f(a : 'IT[:,:,:]', b : 'IT2[:,:,:]'):
+        return a | b
+    def g(a : 'IT[:,:,:]', b : 'IT2'):
+        return a | b
+    def h(a : 'IT', b : 'IT2[:,:,:]'):
+        return a | b
+
+    epyc_f = epyccel(f, language=language)
+    epyc_g = epyccel(g, language=language)
+    epyc_h = epyccel(h, language=language)
+
+    for t_x in int_types:
+        for t_y in int_types:
+            x = randint(2 if t_x is bool else 127, size=(2,3,4), dtype = t_x)
+            y = randint(2 if t_y is bool else 127, size=(2,3,4), dtype = t_y)
+
+            z = f(x,y)
+            epyc_z = epyc_f(x,y)
+            assert np.array_equal(epyc_z, z)
+            assert z.dtype is epyc_z.dtype
+
+            y = randint(2 if t_y is bool else 127, dtype = t_y)
+
+            z2 = g(x,y)
+            epyc_z2 = epyc_g(x,y)
+            assert np.array_equal(epyc_z2, z2)
+            assert z2.dtype is epyc_z2.dtype
+
+            z3 = h(y,x)
+            epyc_z3 = epyc_h(y,x)
+            assert np.array_equal(epyc_z3, z3)
+            assert z3.dtype is epyc_z3.dtype
+
+def test_numpy_bit_xor(language):
+    def f(a : 'IT[:,:,:]', b : 'IT2[:,:,:]'):
+        return a ^ b
+    def g(a : 'IT[:,:,:]', b : 'IT2'):
+        return a ^ b
+    def h(a : 'IT', b : 'IT2[:,:,:]'):
+        return a ^ b
+
+    epyc_f = epyccel(f, language=language)
+    epyc_g = epyccel(g, language=language)
+    epyc_h = epyccel(h, language=language)
+
+    for t_x in int_types:
+        for t_y in int_types:
+            x = randint(2 if t_x is bool else 127, size=(2,3,4), dtype = t_x)
+            y = randint(2 if t_y is bool else 127, size=(2,3,4), dtype = t_y)
+
+            z = f(x,y)
+            epyc_z = epyc_f(x,y)
+            assert np.array_equal(epyc_z, z)
+            assert z.dtype is epyc_z.dtype
+
+            y = randint(2 if t_y is bool else 127, dtype = t_y)
+
+            z2 = g(x,y)
+            epyc_z2 = epyc_g(x,y)
+            assert np.array_equal(epyc_z2, z2)
+            assert z2.dtype is epyc_z2.dtype
+
+            z3 = h(y,x)
+            epyc_z3 = epyc_h(y,x)
+            assert np.array_equal(epyc_z3, z3)
+            assert z3.dtype is epyc_z3.dtype
+
+def test_numpy_bit_lshift(language):
+    def f(a : 'IT[:,:,:]', b : 'IT2[:,:,:]'):
+        return a << b
+    def g(a : 'IT[:,:,:]', b : 'IT2'):
+        return a << b
+    def h(a : 'IT', b : 'IT2[:,:,:]'):
+        return a << b
+
+    epyc_f = epyccel(f, language=language)
+    epyc_g = epyccel(g, language=language)
+    epyc_h = epyccel(h, language=language)
+
+    for t_x in int_types:
+        for t_y in int_types:
+            x = randint(2 if t_x is bool else 32, size=(2,3,4), dtype = t_x)
+            y = randint(2 if t_y is bool else 5, size=(2,3,4), dtype = t_y)
+
+            z = f(x,y)
+            epyc_z = epyc_f(x,y)
+            assert np.array_equal(epyc_z, z)
+            assert z.dtype is epyc_z.dtype
+
+            y = randint(2 if t_y is bool else 5, dtype = t_y)
+
+            z2 = g(x,y)
+            epyc_z2 = epyc_g(x,y)
+            assert np.array_equal(epyc_z2, z2)
+            assert z2.dtype is epyc_z2.dtype
+
+            x = randint(2 if t_x is bool else 32, dtype = t_x)
+            y = randint(2 if t_y is bool else 5, size=(2,3,4), dtype = t_y)
+
+            z3 = h(x,y)
+            epyc_z3 = epyc_h(x,y)
+            assert np.array_equal(epyc_z3, z3)
+            assert z3.dtype is epyc_z3.dtype
+
+def test_numpy_bit_rshift(language):
+    def f(a : 'IT[:,:,:]', b : 'IT2[:,:,:]'):
+        return a >> b
+    def g(a : 'IT[:,:,:]', b : 'IT2'):
+        return a >> b
+    def h(a : 'IT', b : 'IT2[:,:,:]'):
+        return a >> b
+
+    epyc_f = epyccel(f, language=language)
+    epyc_g = epyccel(g, language=language)
+    epyc_h = epyccel(h, language=language)
+
+    for t_x in int_types:
+        for t_y in int_types:
+            x = randint(2 if t_x is bool else 32, size=(2,3,4), dtype = t_x)
+            y = randint(2 if t_y is bool else 5, size=(2,3,4), dtype = t_y)
+
+            z = f(x,y)
+            epyc_z = epyc_f(x,y)
+            assert np.array_equal(epyc_z, z)
+            assert z.dtype is epyc_z.dtype
+
+            y = randint(2 if t_y is bool else 5, dtype = t_y)
+
+            z2 = g(x,y)
+            epyc_z2 = epyc_g(x,y)
+            assert np.array_equal(epyc_z2, z2)
+            assert z2.dtype is epyc_z2.dtype
+
+            x = randint(2 if t_x is bool else 32, dtype = t_x)
+            y = randint(2 if t_y is bool else 5, size=(2,3,4), dtype = t_y)
+
+            z3 = h(x,y)
+            epyc_z3 = epyc_h(x,y)
+            assert np.array_equal(epyc_z3, z3)
+            assert z3.dtype is epyc_z3.dtype
+
+@pytest.mark.skipif(os.environ.get('PYCCEL_DEFAULT_COMPILER', 'GNU') == 'intel',
+        reason="Intel's invert implementation does not seem to match.")
+def test_numpy_bit_invert(language):
+    def f(a : 'IT[:,:,:]'):
+        return ~a
+
+    epyc_f = epyccel(f, language=language)
+
+    for t_x in int_types:
+        x = randint(2 if t_x is bool else 127, size=(2,3,4), dtype = t_x)
+
+        z = f(x)
+        epyc_z = epyc_f(x)
+        assert np.array_equal(epyc_z, z)
+        assert z.dtype is epyc_z.dtype