[REF-2821]Improve Dynamic Imports

[ElijahAhianyo]

This Pr is a big refactor of reflex imports. Much time is spent making imports when components are accessed in reflex. The motivation is to load imports lazily or dynamically only when needed and do this once.

Dynamic Imports

Reflex utilizes dynamic imports, or lazy loading, to reduce startup/import times.
With this approach, imports are delayed until they are needed. We use
the lazy_loader library to achieve this.

How it works

lazy_loader.attach takes two optional arguments: submodules and submod_attrs.

• submodules typically points to directories or files to be accessed.
• submod_attrs defines a mapping of directory or file names as keys with a list
  of attributes or modules to access.

Example directory structure:

reflex/
    |_ components/
            |_ radix/
                |_ themes/
                    |_ components/
                        |_ box.py

To add box under the rx namespace (rx.box), add the relative path to submod_attrs in
reflex/__init__.py (this file):

lazy_loader.attach(
    submodules={"components"},
    submod_attrs={
        "components.radix.themes.components.box": ["box"]
    }
)

This implies that box will be imported from reflex/components/radix/themes/components/box.py.

To add box under the rx.radix namespace (rx.radix.box), add the relative path to the
submod_attrs argument in reflex/components/radix/__init__.py:

lazy_loader.attach(
    submodules = {"themes"},
    submod_attrs = {
        "themes.components.box": ["box"]
    }
)

Note: It is important to specify the immediate submodules of a directory in the submodules
argument to ensure they are registered at runtime. For example, 'components' for reflex,
'radix' for components, 'themes' for radix, etc.

Pyi_generator

To generate .pyi files for __init__.py files, we read the _SUBMODULES and _SUBMOD_ATTRS
attributes to generate the import statements. It is highly recommended to define these with
the provided annotations to facilitate their generation.

Aliases

This is a special case to specify an alias for a component.
As an example, we use this typically for rx.list where defining list attribute in the list.py
overshadows python's list object which messes up the pyi generation for list.pyi. As a result, aliases
should be used for similar cases like this. Note that this logic is employed to fix the pyi generation and alias
should still be defined or accessible. Check out the getattr logic in reflex/components/radix/themes/layouts/list.py

lazy_loader.attach(
    submodules={"components"},
    submod_attrs={
        "components.radix.themes.layouts": [("list_ns", "list")]
    }
)

In the example above, you will be able to do rx.list

Benchmarking

The following reflex code was tested on mac and windows using hyperfine:

import reflex as rx

class State(rx.State):
    """The app state."""
    ...


def index() -> rx.Component:
    # Welcome Page (Index)
    return rx.container(
        rx.color_mode.button(position="top-right"),
        rx.vstack(
            rx.heading("Welcome to Reflex!", size="9"),
            rx.text(
                "Get started by editing ",
                size="5",
            ),
            rx.link(
                rx.button("Check out our docs!"),
                href="https://reflex.dev/docs/getting-started/introduction/",
                is_external=True,
            ),
            spacing="5",
            justify="center",
            min_height="85vh",
        ),
        rx.logo(),
    )


app = rx.App()
app.add_page(index)

Mac

main branch

This PR

Windows

main branch

This PR:

[picklelo]

Based on this: https://pypi.org/project/lazy-loader/ it looks like lazy_loader is only recommended for Python 3.11+ - is it okay to use it for our case?

[ElijahAhianyo]

Based on this: https://pypi.org/project/lazy-loader/ it looks like lazy_loader is only recommended for Python 3.11+ - is it okay to use it for our case?

I think for our case it is, lazy_supports python >= 3.8, however they added the recommendation in their latest release due to a bug in python which affects the use of lazy_loader.load(scientific-python/lazy-loader#88) API in multithreaded env. We're more interested in the lazy_loader.attach API which I think should be fine.

[ElijahAhianyo]

Integration tests need reflex-dev/reflex-web#722 to pass for reflex-web

[linear]

REF-2821 Improve Dynamic import in reflex

[picklelo]

Nice! Changes look good to me, I see a big performance improvement locally when testing.

[Lendemor]

Is that special case still needed despite using the tuple in __init__.py to set the alias of list_ns to list?

[ElijahAhianyo]

The idea of having of having aliases in tuples came about primarily as a result of the list component which shadows python’s list object when explicitly defined. This messed up the pyi generator as well. This was a fix to generate pyi files properly for cases similar to that of rx.list. However, list still needs to be resolved at runtime (reflex.components.radix.themes.layout.list.list) and in reflex.utils.lazy_loader.attach, we only pick the alias than the actual ref

To explain this further:

The mapping in relex.__init__.py contains:

	"components.radix.themes.layouts.list": [("list_ns", "list")]
}

In list.py, assuming list_ns is defined and list is not, the submod_attrs arg of lazy_loader.attach does not accept tuples as list values (only accepts dict[str, list[str]]).
In reflex.utils.lazy_loader we flatten the submod_list as a result keeping the alias(list) leaving out the actual ref (list_ns) as a result, this is also because in lazy_loader.attach implementation, they only dynamically import the module when it(list) exists in the submod_attrs keys(which is recomputed). The issue here is when rx.list is being resolved, the name list instead of list_ns will be passed as the __getattrs__ name argument which will not be found if we added list_ns instead of list in submod_attrs Which means list still needed to be defined to be accessible(https://github.com/scientific-python/lazy_loader/blob/2031223913c7e592c1c59aec72f296e9d08171fa/lazy_loader/__init__.py#L78). 
It may look a bit clunky and confusing(which maybe needs a bit more documentation in the docstrings) considering the “alias” is used primarily for pyi generation but still requires that the attribute be defined as well, but I left it this way since it’s a special case for the list component and didn’t want to over-engineer the approach just yet.

I think another approach will be to have the lazy_loader.attach function implementation defined in reflex.utils.lazy_loader, and then modifying it to work with tuples for __getattr__ and __all__.

[benedikt-bartscher]

@ElijahAhianyo Thanks for the huge performance improvement!

Is this wanted?

Edit: Another finding: AttributeError: No reflex.components.el attribute tbody. Did you mean: 'Tbody'?

[ElijahAhianyo]

@ElijahAhianyo Thanks for the huge performance improvement!

Is this wanted?

Edit: Another finding: AttributeError: No reflex.components.el attribute tbody. Did you mean: 'Tbody'?

@benedikt-bartscher good catch on tbody error. I have a fix for that here.
As for the rx.list, It looks like pylance is treating list as python list object.

Pycharm seem to handle this well