expose_kernel_module.py

'''
Copyright (c) 2010-2023, Delft University of Technology
All rigths reserved

This file is part of the Tudat. Redistribution and use in source and
binary forms, with or without modification, are permitted exclusively
under the terms of the Modified BSD license. You should have received
a copy of the license with this file. If not, please or visit:
http://tudat.tudelft.nl/LICENSE.
'''

# General imports
import os
import re
import inspect
import argparse
from pathlib import Path

#--------------------------------------------------------------------
#%% FUNCTION TO OBTAIN LIST CONTAINING ALL KERNEL MODULES AND SUBMODULES
#--------------------------------------------------------------------

def recursively_find_children(module):
    """
    Recursively find all children of any given module.
    """

    children = inspect.getmembers(module, inspect.ismodule)

    if children:
        for (name, child) in children:
            children += recursively_find_children(child)
        return children
    else:
        return []

#--------------------------------------------------------------------
#%% FUNCTION TO EXPLICITLY IMPORT ALL KERNEL MODULE MEMBERS FOR AUTOCOMPLETION
#--------------------------------------------------------------------

module_methods = lambda module: inspect.getmembers(module, inspect.isroutine)

module_objects = lambda module: [
    (name, value) for name, value in inspect.getmembers(module)
    if (not inspect.isroutine(value) and not inspect.ismodule(value) and not name.startswith('__'))
]

def create_import_statement_to_import_all_module_members(module):

    module_path  = module.__name__

    import_block_title = lambda title: f"# {'-'*40}\n# {title:^40}\n# {'-'*40}"

    # Create import statement
    import_statement = []
    # Import methods
    if module_methods(module):
        import_statement += [
                                import_block_title('METHODS'),
                                f'from {module_path} import \\',
                            ] + [
                                f'     {name}' + (", \\" if i != len(module_methods(module)) - 1 else "") for i, (name, _) in enumerate(module_methods(module))
                            ]
    # Import objects
    if module_objects(module):
        import_statement += [
                                import_block_title('OBJECTS'),
                                f'from {module_path} import \\'
                            ] + [
                                f'     {name}' + (", \\" if i != len(module_objects(module)) - 1 else "") for i, (name, _) in enumerate(module_objects(module))
                            ]
    # Join import statement
    if import_statement:
        import_statement = '\n'.join(import_statement)

    return import_statement


def create_file_to_import_all_module_members(module):

    module_path  = module.__name__
    module_name  = module_path[len('tudatpy.kernel.'):]
    module_name  = module_path[len('tudatpy.kernel.'):]

    # Automatic creation warning
    warning = \
f"""# ====================================================================
#             WARNING: DO NOT MANUALLY CHANGE THIS FILE!
# ====================================================================
# 
# This file is automatically generated before every `tudatpy` release.
# The purpose of this file is manually import all members of 
# 
#     tudatpy.{module_name}
# 
# to make autocompletion suggestions of C++ tudatpy modules possible.
"""

    # Hybrid module path
    hybrid_module_path = Path(f"tudatpy/{module_name.replace('.', '/')}")
    hybrid_module_path.mkdir(parents=True, exist_ok=True)
    kernel_member_import_file_path = hybrid_module_path / '_import_all_kernel_members.py'

    # Create import statement
    import_statement = create_import_statement_to_import_all_module_members(module)

    # Write warning and import statement to `_import_all_kernel_members.py`
    if import_statement:
        with open(kernel_member_import_file_path, 'w') as f:
            f.write(f'{warning}\n{import_statement}')

#--------------------------------------------------------------------
#%% FUNCTION TO SET UP THE __init__.py FILE OF ALL HYBRID MODULES
#--------------------------------------------------------------------

explanation_of_module_exposure_process = lambda module_name: \
f"""# This file, by virtue of the import statement below, merges
# the Tudat kernel module `tudatpy.kernel.{module_name}` with
# its Python extensions defined in `tudatpy/{module_name.replace('.', '/')}`.
# 
# This allows the import of all the C++ and Python submodules of the 
# `{module_name}` kernel module directly from tudatpy:
# 
#     from tudatpy.{module_name} import <any>
# 
# Without the statement below, importing the `{module_name}` kernel module
# would only be possible as follows, and hybrid Python/C++ modules would not
# be posible in tudatpy.
# 
#     from tudatpy.kernel.{module_name} import <any>
# 
# The reason why C++ kernel modules can only be imported as written above 
# is an issue with the `def_submodule` function of pybind11. The issue is discussed
# [here](https://github.com/pybind/pybind11/issues/2639). 
# 
# We circumvent the issue by, for each module and submodule, 
# 
#     1. Creating a Python module (an empty directory with an `__init__.py` file) 
#        inside `tudatpy` with the same name
#     2. Adding the import statement below to the Python module's `__init__.py` 
#        (this file), thereby making the kernel module or submodule from the Python module.
# 
# This workaround was proposed in [this comment](https://github.com/pybind/pybind11/issues/2639#issuecomment-721238757).
# 
# An added benefit of this method is that it makes it possible to write Python extensions 
# and add them to the kernel modules simply by placing them inside the newly created Python 
# module (at the same level as this file), which is not possible with the `def_submodule` 
# function of pybind11.
# An added benefit of this method is that it makes it possible to write Python extensions 
# and add them to the kernel modules simply by placing them inside this module!
"""

def set_up__init__(module):
    """
    Expose all Tudat kernel modules as hybrid C++/Python tudatpy modules.
    """

    module_path  = module.__name__
    module_name  = module_path[len('tudatpy.kernel.'):]
    module_depth = max(len(module_name.split('.')) - 3, 0)

    # Hybrid module path
    hybrid_module_path = Path(f"tudatpy/{module_name.replace('.', '/')}")
    hybrid_module_path.mkdir(parents=True, exist_ok=True)

    # __init__.py
    import_statement = '\n'.join([
                                     # Kernel module import statement
                                     f"from {module_path} import *",
                                 ] + ([
                                          # If there's any module members to import, they will be manually imported in
                                          # `<module_path>/_import_all_kernel_members.py`. If that is the case,
                                          # we import everything in `_import_all_kernel_members.py`.
                                          # This enables autocomplete suggestions for all members of all kernel modules.
                                          f"from tudatpy.{module_name + '.' if module_name else ''}_import_all_kernel_members import *"
                                      ] if create_import_statement_to_import_all_module_members(module) else [])
                                 )
    disclaimer_text  = explanation_of_module_exposure_process(module_name)
    init_file_path   = hybrid_module_path / '__init__.py'

    if not os.path.isfile(init_file_path):
        # If __init__.py does not exist, create it directly
        with open(init_file_path, 'w') as f:
            f.write(f"{disclaimer_text}\n{import_statement}")
    else:
        # Otherwise, if the import statement is not already present in __init__.py,
        # append the explanation of the module exposure process and import statement to __init__.py
        with open(init_file_path, 'r') as f:
            contents = f.read()
        if re.search(re.escape(import_statement), contents):
            # import_statement is already present in __init__.py
            pass
        else:
            with open(init_file_path, 'a') as f:
                f.write(f"\n\n{disclaimer_text}\n{import_statement}")


#--------------------------------------------------------------------
#%% USER INPUT
#--------------------------------------------------------------------

parser = argparse.ArgumentParser(description='Expose hybrid kernel modules.')
parser.add_argument('modules', metavar='module', type=str, nargs='+',
                    help='a list of hybrid kernel modules to expose')
parser.add_argument('--build-dir', metavar='build_dir', type=str,
                    help='the name tudatpy build directory')
parser.add_argument('-i', '--init', action=argparse.BooleanOptionalAction,
                    help='whether to generate the __init__.py files of all kernel modules')
parser.add_argument('-v', '--verbose', action=argparse.BooleanOptionalAction,
                    help='whether to announce the module and submodule being exposed')
args = parser.parse_args()

#--------------------------------------------------------------------
#%% IMPORT COMPILED TUDATPY
#--------------------------------------------------------------------

# Use custom compiled Tudat version
import sys
sys.path.insert(0, f'/mnt/e/studio/professional/work/2023-2024 Tudat/repos/tudat-bundle/{args.build_dir}/tudatpy/')

# Import Tudat kernel
import tudatpy.kernel as kernel

#--------------------------------------------------------------------
#%% KERNEL MODULE EXPOSURE
#--------------------------------------------------------------------

for module_name in args.modules:

    module = getattr(kernel, module_name)

    # Set up __init__.py files for the kernel and its submodules
    if args.init:
        set_up__init__(module)
        for (submodule_name, submodule) in recursively_find_children(module):
            set_up__init__(submodule)

    # Update autocompletion files for the kernel and its submodules
    create_file_to_import_all_module_members(module)
    for (submodule_name, submodule) in recursively_find_children(module):
        # Update autocompletion files
        create_file_to_import_all_module_members(submodule)