__all__-based dynamic exports

init.py

#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2024 Ilya Egorov <[email protected]>
# SPDX-License-Identifier: 0BSD

__all__ = (
    'init',
)

from sys import modules
from pkgutil import iter_modules
from weakref import proxy
from operator import itemgetter
from importlib import import_module
from importlib.util import (
    LazyLoader, find_spec, module_from_spec, resolve_name,
)


def base_import(name, package=None, *, loader_converter=None):
    absolute_name = resolve_name(name, package)
    
    if absolute_name in modules:
        try:
            return modules[absolute_name]
        except KeyError:
            pass
    
    spec = find_spec(absolute_name)
    
    if spec is None:
        raise ModuleNotFoundError(
            f"No module named {absolute_name!r}", name=absolute_name,
        )
    
    if loader_converter is None:
        loader = spec.loader
    else:
        spec.loader = loader = loader_converter(spec.loader)
    
    module = module_from_spec(spec)
    
    actual_module = modules.setdefault(absolute_name, module)
    
    if module is actual_module:
        loader.exec_module(module)
    
    return actual_module


def eager_import(name, package=None):
    return base_import(name, package)


def lazy_import(name, package=None):
    return base_import(name, package, loader_converter=LazyLoader)


def generate_mapping(package, package_content=None, *, lazy=False):
    package_name = None
    package_parent = None
    
    if package_content is None:
        if isinstance(package, str):
            package = import_module(package)
        
        spec = package.__spec__
        
        package_name = spec.name
        package_parent = spec.parent
        
        if package_name != package_parent:
            package = import_module(package_parent)
        
        package_content = dict.fromkeys(
            module_name
            for module_finder, module_name, is_pkg in sorted(iter_modules(
                package.__path__, '.',
            ), key=itemgetter(2))
        )
    
    package_mapping = {}
    
    for module_name, module_dict in package_content.items():
        if package_name is None:
            if isinstance(package, str):
                package = import_module(package)
            
            spec = package.__spec__
            
            package_name = spec.name
            package_parent = spec.parent
        
        absolute_name = resolve_name(module_name, package_parent)
        
        if absolute_name != package_name:
            if module_dict is None:
                module = eager_import(absolute_name)
                
                try:
                    module_dict = dict.fromkeys(module.__all__)
                except AttributeError:
                    module_dict = dict.fromkeys(
                        name
                        for name in dir(module)
                        if not name.startswith('_')
                    )
            elif lazy:
                module = lazy_import(absolute_name)
            else:
                module = eager_import(absolute_name)
            
            for key, name in module_dict.items():
                if name is None:
                    name = key
                
                package_mapping[name] = (module, key)
    
    return package_mapping


def init(package, package_content=None, *, lazy=False):
    if isinstance(package, str):
        package = import_module(package)
    
    package_mapping = generate_mapping(package, package_content, lazy=lazy)
    
    package = proxy(package)
    
    __all__ = sorted(package_mapping)
    
    def __dir__():
        return sorted({*package.__dict__, *__all__})
    
    def __getattr__(name):
        if name in package_mapping:
            module, key = package_mapping[name]
        else:
            raise AttributeError(name)
        
        return package.__dict__.setdefault(name, getattr(module, key))
    
    if lazy:
        package_namespace = {
            '__all__': __all__,
            '__dir__': __dir__,
            '__getattr__': __getattr__,
        }
    else:
        package_namespace = {
            '__all__': __all__,
            '__dir__': __dir__,
            '__getattr__': __getattr__,
            **{
                name: getattr(module, key)
                for name, (module, key) in package_mapping.items()
            },
        }
    
    return package_namespace