From 85df707d636511013d79ee9b561ad414005fbf20 Mon Sep 17 00:00:00 2001
From: relikd <info@relikd.de>
Date: Tue, 20 Dec 2022 00:11:51 +0100
Subject: [PATCH] refactor: yield GroupBySource instead of slugified key

---
 lektor_groupby/groupby.py | 11 ++++++++++-
 lektor_groupby/watcher.py | 20 ++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/lektor_groupby/groupby.py b/lektor_groupby/groupby.py
index b9c9315..aa3806b 100644
--- a/lektor_groupby/groupby.py
+++ b/lektor_groupby/groupby.py
@@ -23,16 +23,21 @@ class GroupBy:
         self._building = False
         self._watcher = []  # type: List[Watcher]
         self._results = []  # type: List[GroupBySource]
+        self._pre_build_priority = []  # type: List[str]  # config.key
         self.resolver = resolver
 
     @property
     def isBuilding(self) -> bool:
         return self._building
 
-    def add_watcher(self, key: str, config: 'AnyConfig') -> Watcher:
+    def add_watcher(
+        self, key: str, config: 'AnyConfig', *, pre_build: bool = False
+    ) -> Watcher:
         ''' Init Config and add to watch list. '''
         w = Watcher(Config.from_any(key, config))
         self._watcher.append(w)
+        if pre_build:
+            self._pre_build_priority.append(w.config.key)
         return w
 
     def queue_all(self, builder: 'Builder') -> None:
@@ -56,6 +61,10 @@ class GroupBy:
                 for w in self._watcher:
                     if w.should_process(record):
                         w.remember(record)
+        # build sources which need building before actual lektor build
+        if self._pre_build_priority:
+            self.make_once(self._pre_build_priority)
+            self._pre_build_priority.clear()
 
     def make_once(self, filter_keys: Optional[Iterable[str]] = None) -> None:
         '''
diff --git a/lektor_groupby/watcher.py b/lektor_groupby/watcher.py
index 136e39c..9046a23 100644
--- a/lektor_groupby/watcher.py
+++ b/lektor_groupby/watcher.py
@@ -19,7 +19,7 @@ class GroupByCallbackArgs(NamedTuple):
 
 GroupingCallback = Callable[[GroupByCallbackArgs], Union[
     Iterator[Any],
-    Generator[Any, Optional[str], None],
+    Generator[Any, Optional[GroupBySource], None],
 ]]
 
 
@@ -73,19 +73,19 @@ class Watcher:
                 key_obj = next(_gen)
                 while True:
                     if self.config.key_obj_fn:
-                        slug = self._persist_multiple(args, key_obj)
+                        vobj = self._persist_multiple(args, key_obj)
                     else:
-                        slug = self._persist(args, key_obj)
-                    # return slugified key and continue iteration
+                        vobj = self._persist(args, key_obj)
+                    # return groupby virtual object and continue iteration
                     if isinstance(_gen, Generator) and not _gen.gi_yieldfrom:
-                        key_obj = _gen.send(slug)
+                        key_obj = _gen.send(vobj)
                     else:
                         key_obj = next(_gen)
             except StopIteration:
                 del _gen
 
     def _persist_multiple(self, args: 'GroupByCallbackArgs', obj: Any) \
-            -> Optional[str]:
+            -> Optional[GroupBySource]:
         # if custom key mapping function defined, use that first
         res = self.config.eval_key_obj_fn(on=args.record,
                                           context={'X': obj, 'ARGS': args})
@@ -96,8 +96,8 @@ class Watcher:
         return self._persist(args, res)  # normal & null replacement
 
     def _persist(self, args: 'GroupByCallbackArgs', obj: Any) \
-            -> Optional[str]:
-        ''' Update internal state. Return slugified string. '''
+            -> Optional[GroupBySource]:
+        ''' Update internal state. Return grouping parent. '''
         if not isinstance(obj, (str, bool, int, float)) and obj is not None:
             raise ValueError(
                 'Unsupported groupby yield type for [{}]:'
@@ -113,7 +113,7 @@ class Watcher:
             slug = self.config.slugify(str(obj)) or None
         # if neither custom mapping succeeded, do not process further
         if not slug or obj is None:
-            return slug
+            return None
         # update internal object storage
         alt = args.record.alt
         if slug not in self._state[alt]:
@@ -125,7 +125,7 @@ class Watcher:
         src.append_child(args.record, obj)
         # reverse reference
         VGroups.of(args.record).add(args.key, src)
-        return slug
+        return src
 
     def remember(self, record: 'Record') -> None:
         self._rmmbr.append(record)