lazy_evaluation_example.py

#!/usr/bin/env python
from __future__ import annotations

from functools import cached_property, partial
import time
from typing import *

####################################################################################################################
### this is not lazy python, this is part of the lazyutils code, scroll to the middle for the actual lazy python ###
####################################################################################################################

# you cannot define the same name more than once
# `single_assignment_dict` prevents accidental redefinition of a name
# except to `main` and `rtn`, which are handled specially
class single_assignment_dict(dict):
    def __setitem__(self, key, value):
        if key in self:
            raise Exception(f'redefinition of "{key}", original value was "{self[key]}", new value "{value}"')
        return super().__setitem__(key, value)


class thunk:
    """
    A thunk is a value that is yet to be evaluated.
    """
    def __init__(self, global_scope, local_scope):
        self.__global_scope = global_scope
        self.__local_scope = local_scope

    @staticmethod
    def compile(name, ann_name, ann):
        if not isinstance(ann, str):
            # print('re-compile', ann)
            return ann
        return compile(ann, f'{name}.{ann_name}: {ann}', mode='eval')

    @staticmethod
    def ensure_value(val):
        return val.value if isinstance(val, thunk) else val

    @cached_property
    def value(self):
        """
        A thunk is evaluated by accessing this property
        """
        cmd = self.__code_object
        # print('evaluating: ', repr(self))
        # if isinstance(cmd, str):
        #     print('eval from string', cmd)
        return eval(cmd, self.__global_scope, self.__local_scope)

    @cached_property
    def __code_object(self):
        c_gcmd = self.__global_scope.get('__compiled_annotations__', {}).get(self.name)
        c_lcmd = self.__local_scope.get('__compiled_annotations__', {}).get(self.name)
        # gcmd = self.__global_scope.get('__annotations__', {}).get(self.name)
        # lcmd = self.__local_scope.get('__annotations__', {}).get(self.name)
        gcmd = lcmd = None
        cmd = c_lcmd or lcmd or c_gcmd or gcmd
        assert not isinstance(cmd, str), cmd
        return cmd

    def __str__(self):
        return str(self.value)

    # all operators will need to be overriden, but here we just do the few we actually need
    def __add__(self, other):
        return self.value + other

    def __radd__(self, other):
        return other + self.value

    def __getitem__(self, key):
        return self.value.__getitem__(key)

    def __eq__(self, key):
        return self.value.__eq__(thunk.ensure_value(key))

    def __call__(self, *args, **kwargs):
        return self.value.__call__(*args, **{k: v for k, v in kwargs})

    def __set_name__(self, ctx, name):
        self.name = name

    def __repr__(self):
        # we abuse the co_filename to store the command definition as a string
        cmd = self.__code_object
        if not isinstance(cmd, str):
            cmd = self.__code_object.co_filename
        else:
            cmd = 'str: ' + self.__code_object
        return f'<thunk {self.name}: {cmd}>'


def create_lazy_evaluation_context(global_scope=None, local_scope=None):
    global_scope = global_scope or dict()
    local_scope = local_scope or global_scope
    ctx = lazy_evaluation(
        global_scope=global_scope,
        local_scope=local_scope,
    )
    # if once_dict is not None:
    #     once_dict = once_dict or single_assignment_dict()
    #     ctx.__prepare__ = lambda name, *args, **kwargs: once_dict
    return ctx


def _compile_annotations(name, attrs):
    return {
        ann_name: thunk.compile(
            name,
            f'<{ann_name}>' if ann_name in ['main', 'rtn'] else ann_name,
            ann,
        ) for ann_name, ann in attrs.items()
    }


def _compile_function_annotation(defun_attr):
    ann_dict = defun_attr.__annotations__
    ann_dict = {k: v for k,v in ann_dict.items() if k != 'params'}
    ann_dict = _compile_annotations(defun_attr.__name__, ann_dict)
    ann_dict['params'] = defun_attr.__annotations__.get('params', '')
    return ann_dict


def compile_defun(attrs):
    for attr_name, attr in attrs.items():
        if isinstance(attr, type):
            attr.__annotations__ = _compile_function_annotation(attr)

def _lazy_evaluation(name, bases, attrs, global_scope=None, local_scope=None):
    def parse_params(par: str) -> List[str]:
        """
        parses `params` declaration of a function:

        class Foo:
            params: (x)
        """

        return list(map(str.strip, par.strip().lstrip('(').rstrip(')').split(',')))

    new_attrs = dict(attrs)

    assert '__compiled_annotations__' in attrs
    assert all([all(not isinstance(ann, str) for ann_name, ann in attr.__annotations__.items() if ann_name != 'params') for attr_name, attr in attrs.items() if isinstance(attr, type)])

    for attr_name, attr in attrs.items():
        if isinstance(attr, type):
            par = attr.__annotations__.pop('params', '')
            params = parse_params(par)
            defun_attrs = dict(vars(attr))
            defun_attrs['__compiled_annotations__'] = attr.__annotations__
            new_attrs[attr_name] = defun(
                name=attr.__name__,
                bases=attr.__bases__,
                attrs=defun_attrs,
                params=params,
                global_scope=global_scope,
            )

    for attr_name in attrs['__annotations__']:
        if attr_name == 'main':
            new_local_scope = {
                '__annotations__': {attr_name: attrs['__annotations__'][attr_name]},
            }
            new_attrs[attr_name] = thunk(global_scope, new_local_scope)
        elif attr_name == 'rtn':
            if global_scope is local_scope:
                rtn = thunk.compile(name, '<rtn>', attrs['__annotations__'][attr_name])
                new_local_scope = {
                    '__annotations__': {attr_name: rtn},
                    '__compiled_annotations__': {attr_name: rtn},
                }
            else:
                new_local_scope = local_scope
            new_attrs[attr_name] = thunk(global_scope, new_local_scope)
        else:
            new_attrs[attr_name] = thunk(global_scope, local_scope)

    type(name, bases, new_attrs) # this calls __set_name__() on all defined attributes

    local_scope.update(new_attrs)

    if 'main' in attrs['__annotations__']:
        assert global_scope is local_scope
        main = attrs['__annotations__'].pop('main')
        return thunk.ensure_value(new_attrs['main'])

    elif 'rtn' in attrs['__annotations__']:
        if global_scope is local_scope:
            rtn = attrs['__annotations__'].pop('rtn')
            return lambda: flattenList(new_attrs['rtn'][-1])
        else:
            return new_attrs['rtn']


def lazy_evaluation(global_scope=None, local_scope=None):
    def __inner_lazy_evaluation(name, bases, attrs):
        new_attrs = dict(attrs)
        new_attrs['__compiled_annotations__'] = _compile_annotations(name, new_attrs['__annotations__'])
        # new_attrs['__annotations__'] = _compile_annotations(name, new_attrs['__annotations__'])
        compile_defun(attrs)
        return _lazy_evaluation(
            name,
            bases,
            new_attrs,
            global_scope=global_scope,
            local_scope=local_scope,
        )
    return __inner_lazy_evaluation


def defun(name, bases, attrs, params, global_scope):
    def _defun(*args, **kwargs):
        assert len(params) == len(args), (name, params, args)
        local_scope = {p: v for p, v in zip(params, args)}
        local_scope.update(kwargs)
        cls = _lazy_evaluation(
            name,
            bases,
            new_attrs,
            global_scope=global_scope,
            local_scope=local_scope,
        )
        return cls
    new_attrs = dict(attrs)
    # compile_defun(attrs)
    # new_attrs['__compiled_annotations__'].update(_compile_annotations(name, new_attrs['__annotations__']))
    _defun.__name__ = name
    return _defun


def flatten(cons):
    while cons != nil:
        yield cons[0]
        cons = cons[1]


def flattenList(cons):
    def _flattenList(cons):
        if isinstance(cons, thunk):
            return _flattenList(thunk.ensure_value(cons))
        if isinstance(cons, tuple):
            return list(map(_flattenList, thunk.ensure_value(flatten(cons))))
        return thunk.ensure_value(cons)
    return _flattenList(cons)



############################
### start of lazy python ###
############################

once_dict = single_assignment_dict()


# create an evaluation namespace, in this case, we create a namespace
# that will modify this module's global variables directly
module_context = create_lazy_evaluation_context(globals(), locals())


class SimpleExample(metaclass=module_context):
    __annotations__ = once_dict

    # Lazy assignment in Python uses the `:` colon operator,
    # not the eager `=` operator
    myVar : 5 + forty

    # Notice that lazy assignments don't have to be written in
    # execution order. You can write in whatever order you
    # want and they'll still evaluate correctly.
    # Useful for literate programming as well.
    ten : 10
    forty : ten + thirty
    thirty : 30

    # Names can only be defined once per lazy
    # evaluation context
    # Commenting out the assignment to `forty` in the
    # next line produces a compile-time exception:
    # forty : ten + 2    # raises Exception: redefinition of forty, original value was 40, new value 40

    # `rtn` is a special construct to execute instructions in order of execution,
    # and returns the last value in the "tuple"/function body.
    # Here we print a string and return `myVar`.  Unlike the rest of
    # lazy_evaluation, order is important in a `rtn` construct and it's useful
    # when you need to call something for their side effect.
    rtn : (
        # only called for the side effect of printing when `rtn` is evaluated
        print('simple'),

        # this value will be returned
        myVar,
    )


# now we're back outside lazy_evaluation context, we're back to normal eager
# evaluation

# now these values can also be read outside the lazy_evaluation context
print(forty)  # output: "40"

# calling the "class" immediately executes the body of SimpleExample.rtn and
# returns it returns value
s = SimpleExample()  # output: "simple"
assert type(s) == int
print('s:', s)       # output: "s: 45"

# note that evaluations are not just lazy, repeating execution does not re-run
# the body of SimpleExample.rtn, instead it just returns its cached value
print(SimpleExample()) # output: "45"


class DefiningListsAndFunctions(metaclass=module_context):
    __annotations__ = once_dict

    # Lists are defined as cons-list, similar to Lisp. In eager Python code,
    # this would be equivalent to the list `myVar = [1, 2, 3, 4]`.
    # If you know linked list, cons-list is almost like one.
    myList : (1, (2, (3, (4, nil)))) = List[int]

    # nil is an empty list, and it always terminates a cons-list
    nil : ()


    # Let's define some simple lazy functions
    #
    # Don't be fooled that we used the `class` keyword here, this is actually
    # defining a function with lazy execution semantic, not a class!

    # We need `head` and `tail`, which are going to be very useful when working
    # with cons-list. `head` and `tail` are also sometimes known as `car` and
    # `cdr` in Lisp.
    #
    # Equivalent Python code would look like this:
    #
    #     def head(cons: list):
    #         return cons[0]
    #     def tail(cons: list):
    #         return cons[1:]
    class head:
        # `params` declares the arguments list that a function takes, here we
        # take one parameter named `cons`
        params : (cons)
        rtn : cons[0]

    class tail:
        params : (cons)

        # Note that with cons-list, taking `cons[1]` returns the whole list
        # except the first item, not just the second item
        rtn : cons[1]


    # Let's see a more complex function definition
    #
    # `tmap` is basically like a `map(func, iterables)` in Python, but works
    # with the cons-list instead of regular list and it is lazy
    #
    # Equivalent Python code would look like this (except we cons-list, we
    # don't actually make copies when slicing it):
    #
    #     def tmap(func: Callable, cons: list) -> list:
    #         if cons == []:
    #             return []
    #         else:
    #             val = func(cons[0])
    #             rest = tmap(func, cons[1:])
    #             return [val] + rest
    class tmap:
        params : (func, cons)

        # `rtn` specifies the return value, in this case we simply return a
        # cons of two other values
        rtn : (
            nil if cons == nil else
            (val, rest)
        )

        # Just like in top-level lazy_evaluation context, function bodies don't
        # have to be defined in execution order.
        val : func(head(cons))

        # Also note that we are calling `tmap` again here inside the function
        # body, `tmap` is a recursive function!
        rest : tmap(func, tail(cons))


class NaturalNumbers(metaclass=module_context):
    __annotations__ = once_dict

    # Here, we're defining an infinite list of `range(1, infinity)` which would
    # only be possible in a language with lazy evaluation semantic!
    #
    # Based on Haskell code: (source: https://stackoverflow.com/questions/19749350/haskell-how-to-generate-this-infinite-list)
    #
    #     nats = 1 : map (+1) nats
    #
    nats : (1, tmap(lambda n: n+1, nats))

    # let's just get the first 20 natural numbers, in python
    # this would be like doing nats[:20].
    first20 : (take(20, nats))

    rtn : (
        print('natural'),
        print(five),

        # inside a `rtn`, we can also use `:=` for assignment to be used in
        # later part of the `rtn` body
        six := five + one,
        print('five + one = ', six),

        # flattenList converts a cons-list to a regular list and tries to fully
        # evaluates all lazy thunks in the result
        flattenList(first20),
    )

    # note that natural number starts with 1, while cons-list indexes starts with 0
    one : nth(0, nats)
    five : nth(4, nats)


class Fibonacci(metaclass=module_context):
    __annotations__ = once_dict

    # Another more complex infinite list example, this time we are doing an
    # infinite Fibonacci list
    #
    # Based on Haskell code: (source: https://stackoverflow.com/questions/50101409/fibonacci-infinite-list-in-haskell)
    #
    #     fibs = 0 : 1 : zipWith (+) fibs (tail fibs)
    #
    fibs : (0, (1, zipWith(bundleArgs(int.__add__), fibs, tail(fibs))))

    rtn : (
        print('fib'),
        print(fib10 + five),
        flattenList(fibs5To15),
    )

    indexedFibs : zipWith(lambda n: (n[0], (n[1], nil)), nats, fibs)
    slowFibs : flattenList(sliced(10, 2000, indexedFibs))

    fib10 : nth(10, fibs)

    fibs5To15 : sliced(5, 15, fibs)


class Helpers(metaclass=module_context):
    __annotations__ = once_dict

    # More function definitions, note that functions don't have to be defined
    # in the same lazy_evaluation() context as the one that they are used on.
    #
    # The context class don't even have to be declared in the order that their
    # functions are used, as long as you don't try to evaluate the values
    # before the context class defining those functions are defined.

    class zipWith:
        params : (func, first, second)
        rtn : (
            nil if first == nil else
            nil if second == nil else
            nextItem
        )

        nextItem : (val, rest)
        val : func((first[0], second[0]))
        rest : zipWith(func, first[1], second[1])

    # equivalent to doing an `list[:n]
    class take:
        params : (n, cons)
        rtn : (
            nil if n == 0 else
            nil if cons == nil else
            nextItem
        )

        nextItem : (head(cons), rest)
        rest : take(n-1, tail(cons))

    # equivalent to doing an `list[n:]
    class drop:
        params : (n, cons)
        rtn : (
            cons if n == 0 else
            cons if cons == nil else
            rest
        )
        rest : drop(n-1, tail(cons))

    # equivalent to doing an `list[n]`
    class nth:
        params : (n, cons)
        rtn : compose2(head, drop)(n, cons)

    # equivalent to doing a `list[start:end]`
    class sliced:
        params : (start, end, cons)
        rtn : drop(start, take(end, cons))

    # given a function with multiple arguments, convert them to a function with one tuple-argument
    class bundleArgs:
        params : (func)
        rtn : lambda args, func=func: func(*map(thunk.ensure_value, args))

    # function composition. converts f(g(n)) -> (f . g)(n)
    class compose2:
        params : (f, g)
        rtn : lambda *args, f=f, g=g, **kwargs: f(g(*args, **kwargs))





print('compilation complete')


class MainFunction(metaclass=module_context):
    __annotations__ = once_dict

    # "main" is similar to `rtn, except that `main` code block is executed
    # immediately when the class is defined, so you don't have to call it from
    # eager Python
    main : (
        print(fib10 + five),
        print(myVar),
        print(five),
        print(flattenList(fibs5To15)),
        print(flattenList(take(100, nats))),

        # this does not actually execute slowFibs
        start := time.time(),
        print("slowFibs wasn't evaluated"),
        a := slowFibs,
        print("time: ", time.time() - start),

        print('forces eager evaluation of slowFibs'),
        start := time.time(),
        thunk.ensure_value(slowFibs),
        print('time: ', time.time() - start),
    )

lazyutils.py

#!/usr/bin/env python
from __future__ import annotations

from functools import cached_property, partial
from typing import *


# you cannot define the same name more than once
# `single_assignment_dict` prevents accidental redefinition of a name
# except to `main` and `rtn`, which are handled specially
class single_assignment_dict(dict):
    def __setitem__(self, key, value):
        if key in self:
            raise Exception(f'redefinition of "{key}", original value was "{self[key]}", new value "{value}"')
        return super().__setitem__(key, value)


class thunk:
    """
    A thunk is a value that is yet to be evaluated.
    """
    def __init__(self, global_scope, local_scope):
        self.__global_scope = global_scope
        self.__local_scope = local_scope

    @staticmethod
    def compile(name, ann_name, ann):
        if not isinstance(ann, str):
            # print('re-compile', ann)
            return ann
        return compile(ann, f'{name}.{ann_name}: {ann}', mode='eval')

    @staticmethod
    def ensure_value(val):
        return val.value if isinstance(val, thunk) else val

    @cached_property
    def value(self):
        """
        A thunk is evaluated by accessing this property
        """
        cmd = self.__code_object
        # print('evaluating: ', repr(self))
        # if isinstance(cmd, str):
        #     print('eval from string', cmd)
        return eval(cmd, self.__global_scope, self.__local_scope)

    @cached_property
    def __code_object(self):
        c_gcmd = self.__global_scope.get('__compiled_annotations__', {}).get(self.name)
        c_lcmd = self.__local_scope.get('__compiled_annotations__', {}).get(self.name)
        # gcmd = self.__global_scope.get('__annotations__', {}).get(self.name)
        # lcmd = self.__local_scope.get('__annotations__', {}).get(self.name)
        gcmd = lcmd = None
        cmd = c_lcmd or lcmd or c_gcmd or gcmd
        assert not isinstance(cmd, str), cmd
        return cmd

    def __str__(self):
        return str(self.value)

    # all operators will need to be overriden, but here we just do the few we actually need
    def __add__(self, other):
        return self.value + other

    def __radd__(self, other):
        return other + self.value

    def __getitem__(self, key):
        return self.value.__getitem__(key)

    def __eq__(self, key):
        return self.value.__eq__(thunk.ensure_value(key))

    def __call__(self, *args, **kwargs):
        return self.value.__call__(*args, **{k: v for k, v in kwargs})

    def __set_name__(self, ctx, name):
        self.name = name

    def __repr__(self):
        # we abuse the co_filename to store the command definition as a string
        cmd = self.__code_object
        if not isinstance(cmd, str):
            cmd = self.__code_object.co_filename
        else:
            cmd = 'str: ' + self.__code_object
        return f'<thunk {self.name}: {cmd}>'


def create_lazy_evaluation_context(global_scope=None, local_scope=None):
    global_scope = global_scope or dict()
    local_scope = local_scope or global_scope
    ctx = lazy_evaluation(
        global_scope=global_scope,
        local_scope=local_scope,
    )
    # if once_dict is not None:
    #     once_dict = once_dict or single_assignment_dict()
    #     ctx.__prepare__ = lambda name, *args, **kwargs: once_dict
    return ctx


def _compile_annotations(name, attrs):
    return {
        ann_name: thunk.compile(
            name,
            f'<{ann_name}>' if ann_name in ['main', 'rtn'] else ann_name,
            ann,
        ) for ann_name, ann in attrs.items()
    }


def _compile_function_annotation(defun_attr):
    ann_dict = defun_attr.__annotations__
    ann_dict = {k: v for k,v in ann_dict.items() if k != 'params'}
    ann_dict = _compile_annotations(defun_attr.__name__, ann_dict)
    ann_dict['params'] = defun_attr.__annotations__.get('params', '')
    return ann_dict


def compile_defun(attrs):
    for attr_name, attr in attrs.items():
        if isinstance(attr, type):
            attr.__annotations__ = _compile_function_annotation(attr)

def _lazy_evaluation(name, bases, attrs, global_scope=None, local_scope=None):
    def parse_params(par: str) -> List[str]:
        """
        parses `params` declaration of a function:

        class Foo:
            params: (x)
        """

        return list(map(str.strip, par.strip().lstrip('(').rstrip(')').split(',')))

    new_attrs = dict(attrs)

    assert '__compiled_annotations__' in attrs
    assert all([all(not isinstance(ann, str) for ann_name, ann in attr.__annotations__.items() if ann_name != 'params') for attr_name, attr in attrs.items() if isinstance(attr, type)])

    for attr_name, attr in attrs.items():
        if isinstance(attr, type):
            par = attr.__annotations__.pop('params', '')
            params = parse_params(par)
            defun_attrs = dict(vars(attr))
            defun_attrs['__compiled_annotations__'] = attr.__annotations__
            new_attrs[attr_name] = defun(
                name=attr.__name__,
                bases=attr.__bases__,
                attrs=defun_attrs,
                params=params,
                global_scope=global_scope,
            )

    for attr_name in attrs['__annotations__']:
        if attr_name == 'main':
            new_local_scope = {
                '__annotations__': {attr_name: attrs['__annotations__'][attr_name]},
            }
            new_attrs[attr_name] = thunk(global_scope, new_local_scope)
        elif attr_name == 'rtn':
            if global_scope is local_scope:
                rtn = thunk.compile(name, '<rtn>', attrs['__annotations__'][attr_name])
                new_local_scope = {
                    '__annotations__': {attr_name: rtn},
                    '__compiled_annotations__': {attr_name: rtn},
                }
            else:
                new_local_scope = local_scope
            new_attrs[attr_name] = thunk(global_scope, new_local_scope)
        else:
            new_attrs[attr_name] = thunk(global_scope, local_scope)

    type(name, bases, new_attrs) # this calls __set_name__() on all defined attributes

    local_scope.update(new_attrs)

    if 'main' in attrs['__annotations__']:
        assert global_scope is local_scope
        main = attrs['__annotations__'].pop('main')
        return thunk.ensure_value(new_attrs['main'])

    elif 'rtn' in attrs['__annotations__']:
        if global_scope is local_scope:
            rtn = attrs['__annotations__'].pop('rtn')
            return lambda: flattenList(new_attrs['rtn'][-1])
        else:
            return new_attrs['rtn']


def lazy_evaluation(global_scope=None, local_scope=None):
    def __inner_lazy_evaluation(name, bases, attrs):
        new_attrs = dict(attrs)
        new_attrs['__compiled_annotations__'] = _compile_annotations(name, new_attrs['__annotations__'])
        # new_attrs['__annotations__'] = _compile_annotations(name, new_attrs['__annotations__'])
        compile_defun(attrs)
        return _lazy_evaluation(
            name,
            bases,
            new_attrs,
            global_scope=global_scope,
            local_scope=local_scope,
        )
    return __inner_lazy_evaluation


def defun(name, bases, attrs, params, global_scope):
    def _defun(*args, **kwargs):
        assert len(params) == len(args), (name, params, args)
        local_scope = {p: v for p, v in zip(params, args)}
        local_scope.update(kwargs)
        cls = _lazy_evaluation(
            name,
            bases,
            new_attrs,
            global_scope=global_scope,
            local_scope=local_scope,
        )
        return cls
    new_attrs = dict(attrs)
    # compile_defun(attrs)
    # new_attrs['__compiled_annotations__'].update(_compile_annotations(name, new_attrs['__annotations__']))
    _defun.__name__ = name
    return _defun


def flatten(cons):
    while cons != nil:
        yield cons[0]
        cons = cons[1]


def flattenList(cons):
    def _flattenList(cons):
        if isinstance(cons, thunk):
            return _flattenList(thunk.ensure_value(cons))
        if isinstance(cons, tuple):
            return list(map(_flattenList, thunk.ensure_value(flatten(cons))))
        return thunk.ensure_value(cons)
    return _flattenList(cons)