Allow LazyLibrary to store Trees as well?

This commit is contained in:
Jan Petykiewicz 2023-01-31 15:31:54 -08:00
parent 02da37a890
commit 81171e9b02

View File

@ -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