#2923 Convert logical sympy expressions back to PSyIR.

Author: hiker

src/psyclone/psyir/frontend/sympy_reader.py

@@ -37,10 +37,37 @@
 '''PSyIR frontend to convert a SymPy expression to PSyIR
 '''
 
+from sympy.printing.printer import Printer
 
 from psyclone.psyir.frontend.fortran import FortranReader
 
 
+# pylint: disable=invalid-name
+class FortranPrinter(Printer):
+    '''A helper class that converts Fortran logical operators back
+    to Fortran format (while SymPy has a Fortran printer (fcode) does this
+    as well, it does not handle e.g. Fortran Array expressions (a(2:5)),
+    so we need to use the not-Fortran-aware output to a normal string,
+    but handle logical operators separately.
+    '''
+    def _print_And(self, expr):
+        '''Called when converting an AND expression.'''
+        return f"({'.AND.' .join(self._print(i) for i in expr.args)})"
+
+    def _print_Or(self, expr):
+        '''Called when converting an OR expression.'''
+        return f"({'.OR.' .join(self._print(i) for i in expr.args)})"
+
+    def _print_Equivalent(self, expr):
+        '''Called when converting an EQUIVALENT expression.'''
+        return f"({'.EQV.' .join(self._print(i) for i in expr.args)})"
+
+    def _print_Xor(self, expr):
+        '''Called when converting an XOR expression, which in Fortran
+        is NEQV.'''
+        return f"({'.NEQV.' .join(self._print(i) for i in expr.args)})"
+
+
 class SymPyReader():
     '''This class converts a SymPy expression, that was created by the
     SymPyWriter, back to PSyIR. It basically allows to use SymPy to modify
@@ -121,7 +148,9 @@ def psyir_from_expression(self, sympy_expr, symbol_table):
         '''
         # Convert the new sympy expression to PSyIR
         reader = FortranReader()
-        return reader.psyir_from_expression(str(sympy_expr), symbol_table)
+        fp = FortranPrinter()
+        return reader.psyir_from_expression(fp.doprint(sympy_expr),
+                                            symbol_table)
 
     # -------------------------------------------------------------------------
     # pylint: disable=no-self-argument, too-many-branches

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

@@ -73,6 +73,18 @@ def test_sympy_reader_constructor():
                                          ("b(:5:2)", "b(:5:2)"),
                                          ("b(2:5:1)", "b(2:5)"),
                                          ("b(2:5:2)", "b(2:5:2)"),
+                                         ("i .and. j", "i .AND. j"),
+                                         ("i .and. j .and. k",
+                                          "i .AND. j .AND. k"),
+                                         ("i .or. j", "i .OR. j"),
+                                         # Precedence requires the ()
+                                         ("i .and. (i .or. j)",
+                                          "i .AND. (i .OR. j)"),
+                                         # Precedence rules discard the ()
+                                         ("i .or. (i .and. j)",
+                                          "i .OR. i .AND. j"),
+                                         ("i .eqv. j", "i .EQV. j"),
+                                         ("i .neqv. j", "i .NEQV. j"),
                                          ])
 def test_sympy_psyir_from_expression(fortran_reader, fortran_writer,
                                      expressions):
@@ -87,7 +99,7 @@ def test_sympy_psyir_from_expression(fortran_reader, fortran_writer,
     '''
     source = f'''program test_prog
                 use my_mod
-                integer :: i, j
+                integer :: i, j, k
                 integer :: a, b(10), c(10, 10)
                 type(my_mod_type) :: d, e(10)
                 x = {expressions[0]}