diff --git a/masque/library.py b/masque/library.py index f25c032..6e84d77 100644 --- a/masque/library.py +++ b/masque/library.py @@ -841,7 +841,7 @@ class LazyLibrary(MutableLibrary): This class is usually used to create a library of Patterns by mapping names to functions which generate or load the relevant `Pattern` object as-needed. """ - dict: Dict[str, Callable[[], 'Pattern']] + dict: Dict[str, Union[Callable[[], 'Pattern'], Callable[[], 'Tree']]] cache: Dict[str, 'Pattern'] _lookups_in_progress: Set[str] @@ -853,7 +853,7 @@ class LazyLibrary(MutableLibrary): def __setitem__( self, key: str, - value: Union['Pattern', Callable[[], 'Pattern']], + value: Union['Pattern', Callable[[], 'Pattern'], Callable[[], 'Tree']], ) -> None: if key in self.dict: raise LibraryError(f'"{key}" already exists in the library. Overwriting is not allowed!') @@ -887,8 +887,14 @@ class LazyLibrary(MutableLibrary): self._lookups_in_progress.add(key) func = self.dict[key] - pat = func() + pat_or_tree = func() self._lookups_in_progress.remove(key) + if isinstance(pat_or_tree, Tree): + tree = pat_or_tree + pat = tree[tree.top] + del tree[tree.top] + self.add(tree) + self.cache[key] = pat return pat