Lazy Val

lazy_val.py

from collections import defaultdict
from typing import Callable


class LazyVal:
    _uncomputed_val = object()

    def __init__(self, fn, *args, **kwargs) -> None:
        self._val = self._uncomputed_val
        self.args = args
        self.kwargs = kwargs
        self._args = list(args)
        self._kwargs = kwargs
        self.fn = fn

    def __repr__(self) -> str:
        if self.is_computed():
            return f"{self.__class__.__name__}({self._val!r})"
        all_args = [f"{v!r}" for v in self._args] + [
            f"{k}={v!r}" for k, v in self._kwargs.items()
        ]
        all_args = ", ".join(all_args)
        return f"{self.__class__.__name__}[{self.fn.__name__}({all_args})]"

    def is_computed(self):
        return self._val is not self._uncomputed_val

    def reset(self):
        self._val = self._uncomputed_val
        self._args = list(self.args)
        self._kwargs = dict(self.kwargs)

    def get(self, recompute=False):
        if not recompute and self.is_computed():
            return self._val
        if recompute:
            self.reset()
            print(f"Recomputing {self}")
        else:
            print(f"Computing {self}")
        for i, arg in enumerate(self._args):
            if isinstance(arg, LazyVal):
                arg = arg.get(recompute=recompute)
            self._args[i] = arg
        for k, arg in self._kwargs.items():
            if isinstance(arg, LazyVal):
                arg = arg.get(recompute=recompute)
            self._kwargs[k] = arg
        self._val = self.fn(*self._args, **self._kwargs)
        return self._val


class PreComputedLazyVal(LazyVal):
    def __init__(self, val, *args, **kwargs) -> None:
        super().__init__(None)
        self._val = val

    def reset(self):
        pass

    def get(self, *args, **kwargs):
        return self._val

if __name__ == "__main__":
    nodes = {}

    def fn_a():
        return "a"

    def fn_b(a):
        return f"b > {a}"

    def fn_c(a, b, d, e=None):
        return f"c -> {{{a}, {b}, {d}, {e=}}}"

    nodes["a"] = LazyVal(fn_a)
    nodes["b"] = LazyVal(fn_b, nodes["a"])
    nodes["d"] = PreComputedLazyVal("d")
    nodes["c"] = LazyVal(fn_c, nodes["a"], nodes["b"], nodes["d"], e="e")

    print(nodes)
    for k, v in nodes.items():
        print(f"Processing {k=}, {v=}")
        print(f"{k=}, {v.get()=}")
        print(nodes)

    print(nodes)
    print("=====" * 10)
    for k, v in reversed(nodes.items()):
        print(f"Processing {k=}, {v=}")
        print(f"{k=}, {v.get(recompute=True)=}")
        print(nodes)

        
"""
Output:

{'a': LazyVal[fn_a()], 'b': LazyVal[fn_b(LazyVal[fn_a()])], 'd': PreComputedLazyVal('d'), 'c': LazyVal[fn_c(LazyVal[fn_a()], LazyVal[fn_b(LazyVal[fn_a()])], PreComputedLazyVal('d'), e='e')]}
Processing k='a', v=LazyVal[fn_a()]
Computing LazyVal[fn_a()]
k='a', v.get()='a'
{'a': LazyVal('a'), 'b': LazyVal[fn_b(LazyVal('a'))], 'd': PreComputedLazyVal('d'), 'c': LazyVal[fn_c(LazyVal('a'), LazyVal[fn_b(LazyVal('a'))], PreComputedLazyVal('d'), e='e')]}
Processing k='b', v=LazyVal[fn_b(LazyVal('a'))]
Computing LazyVal[fn_b(LazyVal('a'))]
k='b', v.get()='b > a'
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal[fn_c(LazyVal('a'), LazyVal('b > a'), PreComputedLazyVal('d'), 
e='e')]}
Processing k='d', v=PreComputedLazyVal('d')
k='d', v.get()='d'
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal[fn_c(LazyVal('a'), LazyVal('b > a'), PreComputedLazyVal('d'), 
e='e')]}
Processing k='c', v=LazyVal[fn_c(LazyVal('a'), LazyVal('b > a'), PreComputedLazyVal('d'), e='e')]
Computing LazyVal[fn_c(LazyVal('a'), LazyVal('b > a'), PreComputedLazyVal('d'), e='e')]
k='c', v.get()="c -> {a, b > a, d, e='e'}"
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
==================================================
Processing k='c', v=LazyVal("c -> {a, b > a, d, e='e'}")
Recomputing LazyVal[fn_c(LazyVal('a'), LazyVal('b > a'), PreComputedLazyVal('d'), e='e')]
Recomputing LazyVal[fn_a()]
Recomputing LazyVal[fn_b(LazyVal('a'))]
Recomputing LazyVal[fn_a()]
k='c', v.get(recompute=True)="c -> {a, b > a, d, e='e'}"
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
Processing k='d', v=PreComputedLazyVal('d')
k='d', v.get(recompute=True)='d'
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
Processing k='b', v=LazyVal('b > a')
Recomputing LazyVal[fn_b(LazyVal('a'))]
Recomputing LazyVal[fn_a()]
k='b', v.get(recompute=True)='b > a'
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
Processing k='a', v=LazyVal('a')
Recomputing LazyVal[fn_a()]
k='a', v.get(recompute=True)='a'
{'a': LazyVal('a'), 'b': LazyVal('b > a'), 'd': PreComputedLazyVal('d'), 'c': LazyVal("c -> {a, b > a, d, e='e'}")}
"""