Fix removal of test_node when rhs.current_user_node becomes PyccelAssociativeParenthesis (pyccel#2140)

This PR fixes an issue where attempting to remove `test_node` from `rhs`
fails when `rhs.current_user_node` becomes a
`PyccelAssociativeParenthesis` during the `operator(lhs, rhs)`
operation. This happens when `rhs` contains more than one element, the
`operator(lhs, rhs)` call creates a `PyccelAssociativeParenthesis`,
which becomes the `current_user_node` of `rhs`. As a result, a direct
call to `rhs.remove_user_node(test_node)` fails since `test_node` no
longer matches the `current_user_node`.

Author: jalalium

CHANGELOG.md

@@ -101,6 +101,7 @@ All notable changes to this project will be documented in this file.
 -   #2111 : Fix declaration of class attributes with name conflicts using type annotations.
 -   #2115 : Fix integer handling with NumPy 2.0 on Windows.
 -   Fix handling of union `typing.TypeAlias` objects as type hints.
+-   #2141 : Fix error when removing `test_node`
 
 ### Changed
 

pyccel/parser/semantic.py

@@ -3652,7 +3652,12 @@ def _visit_AugAssign(self, expr):
             test_node = None
         if test_node:
             lhs.remove_user_node(test_node, invalidate = False)
-            rhs.remove_user_node(test_node, invalidate = False)
+            if test_node in rhs.get_all_user_nodes():
+                rhs.remove_user_node(test_node, invalidate = False)
+            else:
+                assert isinstance(rhs.current_user_node, PyccelAssociativeParenthesis)
+                mid = rhs.current_user_node
+                rhs.remove_user_node(mid, invalidate=False)
             lhs = self._assign_lhs_variable(expr.lhs, self._infer_type(test_node), test_node,
                     new_expressions, is_augassign = True)
             lhs = self._optional_params.get(lhs, lhs)

tests/epyccel/test_epyccel_generators.py

@@ -198,3 +198,12 @@ def f(a0 : 'int[:]'):
     f_epyc = epyccel(f, language = language)
 
     assert f(x) == f_epyc(x)
+
+def test_sum_with_two_variables(language):
+    def f():
+        x = sum(i-j for i in range(10) for j in range(7))
+        return x
+
+    f_epyc = epyccel(f, language=language)
+
+    assert f() == f_epyc()