From 96951c42efe2bae7c48dd6c6fed58b23fd17d9e6 Mon Sep 17 00:00:00 2001
From: serge-sans-paille <serge.guelton@telecom-bretagne.eu>
Date: Wed, 6 Nov 2024 09:19:25 +0100
Subject: [PATCH] [constant folding] Correctly model restrict_assign in
 constant folding

Fix #2247
---
 pythran/optimizations/constant_folding.py |  5 +++++
 pythran/tests/test_optimizations.py       | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/pythran/optimizations/constant_folding.py b/pythran/optimizations/constant_folding.py
index 5bf1f3e02..a16321735 100644
--- a/pythran/optimizations/constant_folding.py
+++ b/pythran/optimizations/constant_folding.py
@@ -13,6 +13,7 @@
 import gast as ast
 from copy import deepcopy
 import logging
+import numpy
 import sys
 
 logger = logging.getLogger('pythran')
@@ -64,6 +65,10 @@ def is_none(val):
     def make_shape(*args):
         return args
 
+    @staticmethod
+    def restrict_assign(expr, value):
+        numpy.copyto(expr, value)
+
 class BreakLoop(Exception):
     pass
 
diff --git a/pythran/tests/test_optimizations.py b/pythran/tests/test_optimizations.py
index 83ad53891..d853a98c3 100644
--- a/pythran/tests/test_optimizations.py
+++ b/pythran/tests/test_optimizations.py
@@ -12,6 +12,20 @@ def test_constant_fold_nan(self):
         code = "def constant_fold_nan(a): from numpy import nan; a[0] = nan; return a"
         self.run_test(code, [1., 2.], constant_fold_nan=[List[float]])
 
+    def test_constant_fold_restrict_assign(self):
+        code = """
+import numpy as np
+def make_signal():
+    y = np.ones((44100,2))
+    y[:, 1:] = 0. #<--- the issue.
+    return y
+
+def constant_fold_restrict_assign():
+    x = make_signal()
+    return x
+    """
+        self.run_test(code, constant_fold_restrict_assign=[])
+
     def test_constant_fold_subscript(self):
         code = '''
 def aux(n):