add Node.find_all

fplll · May 26, 2018 · a555c16 · a555c16
1 parent e0282c5
commit a555c16
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/fpylll/tools/bkz_stats.py b/src/fpylll/tools/bkz_stats.py
@@ -515,6 +515,22 @@ def find(self, label, raise_keyerror=False):
         else:
             return None
 
+    def find_all(self, label):
+        """
+        Find all nodes labelled ``label``
+
+        :param label: a label
+
+        """
+        r = []
+        if self.label == label:
+            r.append(self)
+        if isinstance(self.label, tuple) and self.label[0] == label:
+            r.append(self)
+        for child in self.children:
+            r.extend(child.find_all(label))
+        return tuple(r)
+
     def merge(self, node):
         """
         Merge tree ``node`` into self.