Implementation of Jetpack Compose like reactive tree structure in Python

compose.py

from contextlib import contextmanager


# Helper methods

def pprint_dict(d):
    return  '[' + ', '.join([f'{k} = {d[k]}' for k in d]) + ']'


def composable_tree():
    ui_tree = []
    stack = []

    def get_nearest_branch():
        this_branch = ui_tree
        for br in stack:
            this_branch = this_branch[br]['children']
        return this_branch

    def leaf(func):
        def inner(*args, **kwargs):
            get_nearest_branch().append({
                    'type': 'leaf',
                    'name': func.__name__,
                    'props': kwargs,
                    'render': lambda props: func(**props)
                })

        return inner

    def branch(func):
        @contextmanager
        def inner(*args, **kwargs):
            this_branch = get_nearest_branch()
            this_branch.append({
                    'type': 'branch',
                    'name': func.__name__,
                    'props': kwargs,
                    'children': [],
                    # 'render': lambda props: func(**props, children=this_branch[-1]['children'])
                })
            this_node = this_branch[-1]
            this_node['render'] = lambda props: func(**props, children=this_node['children'])
            try:
                stack_index = len(this_branch) - 1
                stack.append(stack_index)
                yield
            finally:
                stack.pop()

        return inner

    def hook(func):
        def inner(*args, **kwargs):
            get_nearest_branch().append({
                    'type': func.__name__,
                    'value': None
                })
            return func(**kwargs)

        return inner

    def render(component):
        component()

        for top_nodes in ui_tree:
            p = top_nodes.get('props')
            r = top_nodes.get('render')
            if r: r(props=p)

    @hook
    def state(initial):
        this_node = get_nearest_branch()[-1]
        this_node['value'] = this_node['value'] or initial
        this_node['dirty'] = False

        def getState():
            return this_node['value']

        def setState(v):
            this_node['value'] = v
            this_node['dirty'] = True

        return getState, setState

    return render, branch, leaf, state


render, branch, leaf, state = composable_tree()


### Define core components here ###

@branch
def column(d, h, children):
    print('  '*d + '-', 'column', pprint_dict({'h':h}))
    for c in children:
        p = c.get('props')
        r = c.get('render')
        if r: r(props=p)

@branch
def row(d, w, children):
    print('  '*d + '-', 'row', pprint_dict({'w':w}))
    for c in children:
        p = c.get('props')
        r = c.get('render')
        if r: r(props=p)

@leaf
def button(d, text):
    print('  '*d + '-', 'button', pprint_dict({'text':text}))

@leaf
def label(d, text):
    print('  '*d + '-', 'label', pprint_dict({'text':text}))


### Define composite components here ###

@leaf
def header(text):
    print('='*20,text,'='*20)

def item(d,k):
    id, set_id = state(initial="world")
    with row(d=d, w=5):
        button(d=d+1, text=f"Click me x{k} times")
        label(d=d+1, text=f"hello {id()}!")

def app(c):
    header(text='UI')

    for i in range(2):
        with column(d=0, h=4):
            for j in range(2):
                k = i*2+j
                item(1, c)


### Render our UI

counter, set_counter = state(initial=10)
render(lambda: app(counter()))
set_counter(counter()+5)
render(lambda: app(counter()))