Files
lektor-groupby-plugin/lektor_groupby/watcher.py

147 lines
5.5 KiB
Python

from typing import TYPE_CHECKING, Dict, List, Any, Union, NamedTuple
from typing import Optional, Callable, Iterator, Generator
from .backref import VGroups
from .model import ModelReader
from .vobj import GroupBySource
if TYPE_CHECKING:
from lektor.db import Pad, Record
from .config import Config
from .model import FieldKeyPath
class GroupByCallbackArgs(NamedTuple):
record: 'Record'
key: 'FieldKeyPath'
field: Any # lektor model data-field value
GroupingCallback = Callable[[GroupByCallbackArgs], Union[
Iterator[Any],
Generator[Optional[str], Any, None],
]]
class Watcher:
'''
Callback is called with (Record, FieldKeyPath, field-value).
Callback may yield one or more (group, extra-info) tuples.
'''
def __init__(self, config: 'Config') -> None:
self.config = config
self._root = self.config.root
def grouping(self, flatten: bool = True) \
-> Callable[[GroupingCallback], None]:
'''
Decorator to subscribe to attrib-elements.
If flatten = False, dont explode FlowType.
(record, field-key, field) -> (group, extra-info)
'''
def _decorator(fn: GroupingCallback) -> None:
self.flatten = flatten
self.callback = fn
return _decorator
def initialize(self, pad: 'Pad') -> None:
''' Reset internal state. You must initialize before each build! '''
assert callable(self.callback), 'No grouping callback provided.'
self._model_reader = ModelReader(pad.db, self.config.key, self.flatten)
self._root_record = {} # type: Dict[str, Record]
self._state = {} # type: Dict[str, Dict[Optional[str], GroupBySource]]
self._rmmbr = [] # type: List[Record]
for alt in pad.config.iter_alternatives():
self._root_record[alt] = pad.get(self._root, alt=alt)
self._state[alt] = {}
def should_process(self, node: 'Record') -> bool:
''' Check if record path is being watched. '''
return str(node['_path']).startswith(self._root)
def process(self, record: 'Record') -> None:
'''
Will iterate over all record fields and call the callback method.
Each record is guaranteed to be processed only once.
'''
for key, field in self._model_reader.read(record):
args = GroupByCallbackArgs(record, key, field)
_gen = self.callback(args)
try:
group = next(_gen)
while True:
if self.config.key_map_fn:
slug = self._persist_multiple(args, group)
else:
slug = self._persist(args, group)
# return slugified group key and continue iteration
if isinstance(_gen, Generator) and not _gen.gi_yieldfrom:
group = _gen.send(slug)
else:
group = next(_gen)
except StopIteration:
del _gen
def _persist_multiple(self, args: 'GroupByCallbackArgs', obj: Any) \
-> Optional[str]:
# if custom key mapping function defined, use that first
res = self.config.eval_key_map_fn(on=args.record,
context={'X': obj, 'SRC': args})
if isinstance(res, (list, tuple)):
for k in res:
self._persist(args, k) # 1-to-n replacement
return None
return self._persist(args, res) # normal & null replacement
def _persist(self, args: 'GroupByCallbackArgs', obj: Any) \
-> Optional[str]:
''' Update internal state. Return slugified string. '''
if not isinstance(obj, (str, bool, int, float)) and obj is not None:
raise ValueError(
'Unsupported groupby yield type for [{}]:'
' {} (expected str, got {})'.format(
self.config.key, obj, type(obj).__name__))
if obj is None:
# if obj is not set, test if config.replace_none_key is set
slug = self.config.replace_none_key
obj = slug
else:
# if obj is set, apply config.key_map (convert int -> str)
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
# update internal object storage
alt = args.record.alt
if slug not in self._state[alt]:
src = GroupBySource(self._root_record[alt], slug)
self._state[alt][slug] = src
else:
src = self._state[alt][slug]
src.append_child(args.record, obj) # obj is used as "group" string
# reverse reference
VGroups.of(args.record).add(args.key, src)
return slug
def remember(self, record: 'Record') -> None:
self._rmmbr.append(record)
def iter_sources(self) -> Iterator[GroupBySource]:
''' Prepare and yield GroupBySource elements. '''
for x in self._rmmbr:
self.process(x)
del self._rmmbr
for vobj_list in self._state.values():
for vobj in vobj_list.values():
yield vobj.finalize(self.config)
# cleanup. remove this code if you'd like to iter twice
del self._model_reader
del self._root_record
del self._state
def __repr__(self) -> str:
return '<GroupByWatcher key="{}" enabled={}>'.format(
self.config.key, self.config.enabled)