3189 sympy handling equal

src/psyclone/psyir/backend/sympy_writer.py

@@ -111,7 +111,9 @@ class SymPyWriter(FortranWriter):
         {BinaryOperation.Operator.AND: "And({lhs}, {rhs})",
          BinaryOperation.Operator.OR: "Or({lhs}, {rhs})",
          BinaryOperation.Operator.EQV: "Equivalent({lhs}, {rhs})",
-         BinaryOperation.Operator.NEQV: "Xor({lhs}, {rhs})"}
+         BinaryOperation.Operator.NEQV: "Xor({lhs}, {rhs})",
+         BinaryOperation.Operator.EQ: "Eq({lhs}, {rhs})"
+         }
 
     def __init__(self):
         super().__init__()

src/psyclone/psyir/frontend/sympy_reader.py

@@ -49,29 +49,36 @@ class FortranPrinter(Printer):
     not handle e.g. Fortran Array expressions (a(2:5)), so we specialise the
     generic SymPy Printer and handle the necessary conversions.'''
 
-    def _print_And(self, expr):
+    def _print_And(self, expr) -> str:
         '''Called when converting an AND expression.'''
         return f"({'.AND.' .join(self._print(i) for i in expr.args)})"
 
-    def _print_Or(self, expr):
+    def _print_Or(self, expr) -> str:
         '''Called when converting an OR expression.'''
         return f"({'.OR.' .join(self._print(i) for i in expr.args)})"
 
-    def _print_Equivalent(self, expr):
+    def _print_Equivalent(self, expr) -> str:
         '''Called when converting an EQUIVALENT expression.'''
         return f"({'.EQV.' .join(self._print(i) for i in expr.args)})"
 
-    def _print_Xor(self, expr):
+    def _print_Xor(self, expr) -> str:
         '''Called when converting an XOR expression, which in Fortran
         is .NEQV.'''
         return f"({'.NEQV.' .join(self._print(i) for i in expr.args)})"
 
+    def _print_Equality(self, expr) -> str:
+        '''Called when converting an Eq expression, which in Fortran
+        is =='''
+        return f"({'==' .join(self._print(i) for i in expr.args)})"
+
     def _print_BooleanTrue(self, expr) -> str:
         '''Called when converting a SymPy value of True.'''
+        # pylint: disable=unused-argument
         return ".TRUE."
 
     def _print_BooleanFalse(self, expr) -> str:
         '''Called when converting a SymPy value of False.'''
+        # pylint: disable=unused-argument
         return ".FALSE."
 
 

src/psyclone/tests/core/symbolic_maths_test.py

@@ -483,7 +483,11 @@ def test_symbolic_math_use_range(fortran_reader, expressions):
      "norm_u(idx + iw2) * u_e(idx - iw2v + LBOUND(u_e, 1),df2)"),
     (".true. .and. .false.", ".false."),
     ("zh_cum1(jk1) <= zh_cum0(jk0) .AND. zh_cum1(jk1) > zh_cum0(jk0 - 1)",
-     "zh_cum0(jk0) >= zh_cum1(jk1) .AND. zh_cum1(jk1) > zh_cum0(jk0 - 1)")])
+     "zh_cum0(jk0) >= zh_cum1(jk1) .AND. zh_cum1(jk1) > zh_cum0(jk0 - 1)"),
+    ("zh_cum1(jk1) == zh_cum0(jk0) .AND. zh_cum1(jk1) == zh_cum0(jk0 - 1)",
+     "zh_cum0(jk0) == zh_cum1(jk1) .AND. zh_cum1(jk1) == zh_cum0(jk0 - 1)"),
+    ("i == i .and. c(i,i) == 5", "c(i,i) == 5"),
+    ])
 def test_symbolic_maths_expand(fortran_reader, fortran_writer, expr, expected):
     '''Test the expand method works as expected.'''
     # A dummy program to easily create the PSyIR for the

src/psyclone/tests/psyir/backend/sympy_writer_test.py

@@ -43,6 +43,7 @@
 from sympy import Function, Symbol
 from sympy.parsing.sympy_parser import parse_expr
 
+from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.frontend.sympy_reader import SymPyReader
 from psyclone.psyir.backend.sympy_writer import SymPyWriter
 from psyclone.psyir.backend.visitor import VisitorError
@@ -256,7 +257,7 @@ def test_sym_writer_rename_members(fortran_reader, expressions):
 
 
 @pytest.mark.parametrize(
-    "expr, positive, sym_map",
+    "expr, positive, expected_sym_map",
     [("a%x", False, {"a_x": Symbol("a%x")}),
      ("a%x", True, {"a_x": Symbol("a%x", **{"positive": True})}),
      ("a%x(i)", False, {"a_x": Function("a_x"), "i": Symbol("i")}),
@@ -269,11 +270,15 @@ def test_sym_writer_rename_members(fortran_reader, expressions):
                             "b_c": Function("b_c"),
                             "i": Symbol("i", **{"positive": True})}),
      ])
-def test_sym_writer_symbol_types(fortran_reader, expr, positive, sym_map):
+def test_sym_writer_symbol_types(fortran_reader: FortranReader,
+                                 expr: str,
+                                 positive: bool,
+                                 expected_sym_map: dict[str, Symbol]):
     '''Tests that arrays are detected as SymPy functions, and scalars
     as SymPy symbols. The 'expr' parameter contains the expression to parse,
     'positive' is whether or not to flag symbols as positive definite and
-    'sym_map' is the expected mapping of names to SymPy functions or symbols.
+    'expected_sym_map' is the expected mapping of names to SymPy functions
+    or symbols.
 
     '''
     # A dummy program to easily create the PSyIR for the
@@ -286,12 +291,12 @@ def test_sym_writer_symbol_types(fortran_reader, expr, positive, sym_map):
                 end program test_prog '''
 
     psyir = fortran_reader.psyir_from_source(source)
-    expr = psyir.children[0].children[0].rhs
+    expr_psyir = psyir.children[0].children[0].rhs
     sympy_writer = SymPyWriter()
-    _ = sympy_writer(expr, all_variables_positive=positive)
-    assert len(sympy_writer.type_map) == len(sym_map)
-    for key in sympy_writer.type_map.keys():
-        assert sympy_writer.type_map[key] == sym_map[key]
+    _ = sympy_writer(expr_psyir, all_variables_positive=positive)
+    assert len(sympy_writer.type_map) == len(expected_sym_map)
+    for key, sym_map in sympy_writer.type_map.items():
+        assert sym_map == expected_sym_map[key]
 
 
 @pytest.mark.parametrize("expr, sym_map",
@@ -517,14 +522,13 @@ def test_sympy_writer_user_types(fortran_reader, fortran_writer,
                           ("a .or. b", "Or(a, b)"),
                           ("a .eqv. b", "Equivalent(a, b)"),
                           ("a .neqv. b", "Xor(a, b)"),
+                          ("a == b", "Eq(a, b)"),
                           ])
-def test_sympy_writer_logicals(fortran_reader, fortran_writer,
-                               fortran_expr, sympy_str):
-    '''Test handling of user-defined types, e.g. conversion of
-    ``a(i)%b(j)`` to ``a_b(i,i,1,j,j,1)``. Each Fortran expression
-    ``fortran_expr`` is first converted to a string ``sympy_str`` to be
-    parsed by SymPy. The sympy expression is then converted back to PSyIR.
-    This string must be the same as the original ``fortran_expr``.
+def test_sympy_writer_logicals(fortran_reader: FortranReader,
+                               fortran_expr: str,
+                               sympy_str: str):
+    '''Test writing of logical expressions, i.e. that the Fortran
+    logical expressions are correctly converted to SymPy expressions.
 
     '''
     source = f'''program test_prog

src/psyclone/tests/psyir/frontend/sympy_reader_test.py

@@ -37,7 +37,9 @@
 
 import pytest
 
+from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.psyir.backend.sympy_writer import SymPyWriter
+from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.frontend.sympy_reader import SymPyReader
 
 
@@ -88,9 +90,12 @@ def test_sympy_reader_constructor():
                                          (".TRUE. .and. .true.", ".true."),
                                          (".TRUE. .AND. .FALSE.",
                                           ".false."),
+                                         ("b(i) == 3 .and. c(i,i) == 5",
+                                          "b(i) == 3 .AND. c(i,i) == 5"),
                                          ])
-def test_sympy_psyir_from_expression(fortran_reader, fortran_writer,
-                                     expressions):
+def test_sympy_psyir_from_expression(fortran_reader: FortranReader,
+                                     fortran_writer: FortranWriter,
+                                     expressions: tuple[str, str]):
     '''Test conversion from a SymPy expression back to PSyIR. We use the
     SymPyWriter to convert the Fortran string to a SymPy expression (we need
     a symbol table to convert from SymPy to PSyIR, which we get this way