Generalize LDRAW file-tree walking between LdrawDb.tree, .count_studs.

diff --git a/src/bricksplom/ldraw.py b/src/bricksplom/ldraw.py
index c0a22589e8378be52c0afcc2e3aa621818189a6f..d0134cbf8dea34be3b1a53baf832aeff541f17f6 100644 (file)
--- a/src/bricksplom/ldraw.py
+++ b/src/bricksplom/ldraw.py
@@ -2,6 +2,7 @@
 
 # standard libs
 from pathlib import Path
+from typing import Any, Callable
 # ourselves
 from bricksplom.constants import CHAR_NEWLINE
 
@@ -20,6 +21,7 @@ LDRAW_STUDS = {
 
 LdrawRefsTree = tuple[Path, '_SubRefs']
 _SubRefs = dict[str, LdrawRefsTree]
+_Walker = Callable[[str], Any]
 
 
 class LdrawDb:
@@ -42,34 +44,46 @@ class LdrawDb:
                 return path
         raise Exception(f'invalid ldraw reference: {filename}')
 
+    def _walk_tree(
+            self,
+            filename: str,
+            collect: Callable[[_Walker, str, list[str]], Any],
+            result: Callable[[Path, list[Any]], Any],
+            ) -> Any:
+        path = self.path(filename)
+        collected = []
+        for line in path.read_text(encoding='utf8').split(CHAR_NEWLINE):
+            toks = line.lstrip().split(maxsplit=14)
+            if not (toks and toks[0] == '1'):
+                continue
+            assert len(toks) == 15
+            collected += [
+                collect(lambda ref: self._walk_tree(ref, collect, result),
+                        toks.pop().replace('\\', '/'),
+                        toks)
+                ]
+        return result(path, collected)
+
     def tree(
             self,
             filename: str
             ) -> LdrawRefsTree:
         'Produce LdrawRefsTree for filename.'
-        path = self.path(filename)
-        refs: _SubRefs = {}
-        for line in path.read_text(encoding='utf8').split(CHAR_NEWLINE):
-            line = line.lstrip()
-            if not line.startswith('1'):
-                continue
-            ref = line.split(maxsplit=14)[-1].replace('\\', '/')
-            refs[ref] = self.tree(ref)
-        return path, refs
+        return self._walk_tree(
+                filename,
+                collect=lambda walker, ref, _: (ref, walker(ref)),
+                result=lambda path, coll: (path, {t[0]: t[1] for t in coll}))
 
     def count_studs(
             self,
             filename
             ) -> int:
         'On item at {._root}/p(art)/{filename} count studs.'
-        path = self.path(filename)
-        studs_counted = 0
-        for line in path.read_text(encoding='utf8').split(CHAR_NEWLINE):
-            line = line.lstrip()
-            if not line.startswith('1'):
-                continue
-            toks = line.split(maxsplit=14)
-            target = toks[14].replace('\\', '/')
+        def count_studs(
+                walker: _Walker,
+                ref: str,
+                toks: list[str]
+                ) -> int:
             matrix = tuple(tuple(toks[n:n+3]) for n in (5, 8, 11))
             do_ignore = False
             for row in [[float(x) for x in row] for row in matrix]:
@@ -83,6 +97,9 @@ class LdrawDb:
                                       if jdx != idx and row[jdx] != 0]:
                         do_ignore = True
                         break
-            studs_counted += (int((not do_ignore) and target in LDRAW_STUDS)
-                              or self.count_studs(target))
-        return studs_counted
+            return (int((not do_ignore) and ref in LDRAW_STUDS)
+                    or walker(ref))
+
+        return self._walk_tree(filename,
+                               collect=count_studs,
+                               result=lambda _, to_sum: sum(to_sum))