From 7433c5dbb92099b5b4008d776ff9b959b80d5ead Mon Sep 17 00:00:00 2001
From: Cosimo Lupo <clupo@google.com>
Date: Wed, 8 Mar 2023 19:13:48 +0000
Subject: [PATCH] [otTraverse] allow to use custom callback to iterate over
 subtables

---
 Lib/fontTools/ttLib/tables/otTraverse.py | 31 +++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/Lib/fontTools/ttLib/tables/otTraverse.py b/Lib/fontTools/ttLib/tables/otTraverse.py
index 412f89f26..bf22dcfdb 100644
--- a/Lib/fontTools/ttLib/tables/otTraverse.py
+++ b/Lib/fontTools/ttLib/tables/otTraverse.py
@@ -31,6 +31,9 @@ def dfs_base_table(
     root_accessor: Optional[str] = None,
     skip_root: bool = False,
     predicate: Optional[Callable[[SubTablePath], bool]] = None,
+    iter_subtables_fn: Optional[
+        Callable[[BaseTable], Iterable[BaseTable.SubTableEntry]]
+    ] = None,
 ) -> Iterable[SubTablePath]:
     """Depth-first search tree of BaseTables.
 
@@ -43,6 +46,9 @@ def dfs_base_table(
         predicate (Optional[Callable[[SubTablePath], bool]]): function to filter out
             paths. If True, the path is yielded and its subtables are added to the
             queue. If False, the path is skipped and its subtables are not traversed.
+        iter_subtables_fn (Optional[Callable[[BaseTable], Iterable[BaseTable.SubTableEntry]]]):
+            function to iterate over subtables of a table. If None, the default
+            BaseTable.iterSubTables() is used.
 
     Yields:
         SubTablePath: tuples of BaseTable.SubTableEntry(name, table, index) namedtuples
@@ -56,6 +62,7 @@ def dfs_base_table(
         skip_root,
         predicate,
         lambda frontier, new: frontier.extendleft(reversed(new)),
+        iter_subtables_fn,
     )
 
 
@@ -64,11 +71,14 @@ def bfs_base_table(
     root_accessor: Optional[str] = None,
     skip_root: bool = False,
     predicate: Optional[Callable[[SubTablePath], bool]] = None,
+    iter_subtables_fn: Optional[
+        Callable[[BaseTable], Iterable[BaseTable.SubTableEntry]]
+    ] = None,
 ) -> Iterable[SubTablePath]:
     """Breadth-first search tree of BaseTables.
 
     Args:
-        root (BaseTable): the root of the tree.
+    the root of the tree.
         root_accessor (Optional[str]): attribute name for the root table, if any (mostly
             useful for debugging).
         skip_root (Optional[bool]): if True, the root itself is not visited, only its
@@ -76,6 +86,9 @@ def bfs_base_table(
         predicate (Optional[Callable[[SubTablePath], bool]]): function to filter out
             paths. If True, the path is yielded and its subtables are added to the
             queue. If False, the path is skipped and its subtables are not traversed.
+        iter_subtables_fn (Optional[Callable[[BaseTable], Iterable[BaseTable.SubTableEntry]]]):
+            function to iterate over subtables of a table. If None, the default
+            BaseTable.iterSubTables() is used.
 
     Yields:
         SubTablePath: tuples of BaseTable.SubTableEntry(name, table, index) namedtuples
@@ -89,6 +102,7 @@ def bfs_base_table(
         skip_root,
         predicate,
         lambda frontier, new: frontier.extend(new),
+        iter_subtables_fn,
     )
 
 
@@ -98,6 +112,9 @@ def _traverse_ot_data(
     skip_root: bool,
     predicate: Optional[Callable[[SubTablePath], bool]],
     add_to_frontier_fn: AddToFrontierFn,
+    iter_subtables_fn: Optional[
+        Callable[[BaseTable], Iterable[BaseTable.SubTableEntry]]
+    ] = None,
 ) -> Iterable[SubTablePath]:
     # no visited because general otData cannot cycle (forward-offset only)
     if root_accessor is None:
@@ -108,6 +125,11 @@ def _traverse_ot_data(
         def predicate(path):
             return True
 
+    if iter_subtables_fn is None:
+
+        def iter_subtables_fn(table):
+            return table.iterSubTables()
+
     frontier: Deque[SubTablePath] = deque()
 
     root_entry = BaseTable.SubTableEntry(root_accessor, root)
@@ -116,7 +138,10 @@ def _traverse_ot_data(
     else:
         add_to_frontier_fn(
             frontier,
-            [(root_entry, subtable_entry) for subtable_entry in root.iterSubTables()],
+            [
+                (root_entry, subtable_entry)
+                for subtable_entry in iter_subtables_fn(root)
+            ],
         )
 
     while frontier:
@@ -130,7 +155,7 @@ def _traverse_ot_data(
         yield SubTablePath(path)
 
         new_entries = [
-            path + (subtable_entry,) for subtable_entry in current.iterSubTables()
+            path + (subtable_entry,) for subtable_entry in iter_subtables_fn(current)
         ]
 
         add_to_frontier_fn(frontier, new_entries)