From 681d7a7ed3afc66444f60a5fe236c69aafe74608 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Fri, 2 Aug 2019 15:35:41 +0200
Subject: [PATCH] Forward {find,match,subs,replace} of Assignment to
 AssignmentCollection

---
 pystencils/simp/assignment_collection.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py
index b15b87a4f..da4832494 100644
--- a/pystencils/simp/assignment_collection.py
+++ b/pystencils/simp/assignment_collection.py
@@ -362,6 +362,29 @@ class AssignmentCollection:
         self.sub_expressions = [Assignment(k, v)
                                 for k, v in sub_expressions_dict.items()]
 
+    def find(self, *args, **kwargs):
+        return set.union(*[a.find(*args, **kwargs) for a in self.all_assignments])
+
+    def match(self, *args, **kwargs):
+        rtn = {}
+        for a in self.all_assignments:
+            partial_result = a.match(*args, **kwargs)
+            if partial_result:
+                rtn.update(partial_result)
+        return rtn
+
+    def subs(self, *args, **kwargs):
+        return AssignmentCollection(
+            main_assignments=[a.subs(*args, **kwargs) for a in self.main_assignments],
+            subexpressions=[a.subs(*args, **kwargs) for a in self.subexpressions]
+        )
+
+    def replace(self, *args, **kwargs):
+        return AssignmentCollection(
+            main_assignments=[a.replace(*args, **kwargs) for a in self.main_assignments],
+            subexpressions=[a.replace(*args, **kwargs) for a in self.subexpressions]
+        )
+
 
 class SymbolGen:
     """Default symbol generator producing number symbols ζ_0, ζ_1, ..."""
-- 
GitLab