SciCompMod · reneSchm · Jun 10, 2024 · May 28, 2024 · May 28, 2024 · May 29, 2024
diff --git a/.gitignore b/.gitignore
@@ -271,6 +271,9 @@ pycode/epidemiology/**/*.h5
 pycode/epi_venv/
 pycode/build_pylint/
 
+# Automatic generated python code
+pycode/memilio-simulation-stubs/
+
 # Credentials for data download
 pycode/memilio-epidata/memilio/epidata/CredentialsRegio.ini
 

diff --git a/pycode/memilio-simulation/CMakeLists.txt b/pycode/memilio-simulation/CMakeLists.txt
@@ -32,16 +32,30 @@ endif()
 # add in C++ library
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../cpp ${CMAKE_CURRENT_BINARY_DIR}/cpp EXCLUDE_FROM_ALL)
 
+# Function to add a pybind11 module
+function(add_pymio_module target_name)
+    set(options)
+    set(oneValueArgs)
+    set(multiValueArgs LINKED_LIBRARIES SOURCES)
+    cmake_parse_arguments(PYBIND11_MODULE "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+
+    pybind11_add_module(${target_name} MODULE ${PYBIND11_MODULE_SOURCES})
+    target_link_libraries(${target_name} PRIVATE ${PYBIND11_MODULE_LINKED_LIBRARIES})
+    target_include_directories(${target_name} PRIVATE memilio/simulation)
+    install(TARGETS ${target_name} LIBRARY DESTINATION memilio)
+endfunction()
+
+
 # build python extensions
-pybind11_add_module(_simulation_abm MODULE
-    memilio/simulation/abm.cpp
+add_pymio_module(_simulation_abm
+    LINKED_LIBRARIES memilio abm
+    SOURCES memilio/simulation/abm.cpp
 )
-target_link_libraries(_simulation_abm PRIVATE memilio abm)
-target_include_directories(_simulation_abm PRIVATE memilio/simulation)
-install(TARGETS _simulation_abm LIBRARY DESTINATION memilio)
 
-pybind11_add_module(_simulation MODULE
-    memilio/simulation/simulation.cpp
+add_pymio_module(_simulation
+    LINKED_LIBRARIES memilio
+    SOURCES memilio/simulation/simulation.cpp
     memilio/simulation/epidemiology/damping_sampling.cpp
     memilio/simulation/epidemiology/uncertain_matrix.cpp
     memilio/simulation/mobility/metapopulation_mobility_instant.cpp
@@ -51,34 +65,23 @@ pybind11_add_module(_simulation MODULE
     memilio/simulation/utils/parameter_distributions.cpp
     memilio/simulation/utils/uncertain_value.cpp
 )
-target_link_libraries(_simulation PRIVATE memilio)
-target_include_directories(_simulation PRIVATE memilio/simulation)
-install(TARGETS _simulation LIBRARY DESTINATION memilio)
 
-pybind11_add_module(_simulation_osir MODULE
-    memilio/simulation/osir.cpp
+add_pymio_module(_simulation_osir
+    LINKED_LIBRARIES memilio ode_sir
+    SOURCES memilio/simulation/osir.cpp
 )
-target_link_libraries(_simulation_osir PRIVATE memilio ode_sir)
-target_include_directories(_simulation_osir PRIVATE memilio/simulation)
-install(TARGETS _simulation_osir LIBRARY DESTINATION memilio)
 
-pybind11_add_module(_simulation_oseir MODULE
-    memilio/simulation/oseir.cpp
+add_pymio_module(_simulation_oseir
+    LINKED_LIBRARIES memilio ode_seir
+    SOURCES memilio/simulation/oseir.cpp
 )
-target_link_libraries(_simulation_oseir PRIVATE memilio ode_seir)
-target_include_directories(_simulation_oseir PRIVATE memilio/simulation)
-install(TARGETS _simulation_oseir LIBRARY DESTINATION memilio)
 
-pybind11_add_module(_simulation_osecir MODULE
-    memilio/simulation/osecir.cpp
+add_pymio_module(_simulation_osecir
+    LINKED_LIBRARIES memilio ode_secir
+    SOURCES memilio/simulation/osecir.cpp
 )
-target_link_libraries(_simulation_osecir PRIVATE memilio ode_secir)
-target_include_directories(_simulation_osecir PRIVATE memilio/simulation)
-install(TARGETS _simulation_osecir LIBRARY DESTINATION memilio)
 
-pybind11_add_module(_simulation_osecirvvs MODULE
-    memilio/simulation/osecirvvs.cpp
-)
-target_link_libraries(_simulation_osecirvvs PRIVATE memilio ode_secirvvs)
-target_include_directories(_simulation_osecirvvs PRIVATE memilio/simulation)
-install(TARGETS _simulation_osecirvvs LIBRARY DESTINATION memilio)
+add_pymio_module(_simulation_osecirvvs
+    LINKED_LIBRARIES memilio ode_secirvvs
+    SOURCES memilio/simulation/osecirvvs.cpp
+)
diff --git a/pycode/memilio-simulation/README.md b/pycode/memilio-simulation/README.md
@@ -38,6 +38,16 @@ python setup.py install -- -DCMAKE_BUILD_TYPE=Debug -DMEMILIO_USE_BUNDLED_PYBIND
 
 Alternatively, the `CMakeCache.txt` in the directory created by Scikit-Build can be edited to set the options.
 
+## Stubs
+
+A stub file is a file containing a skeleton of the public interface of that Python module including classes, variables, functions and their types. They help by enabling autocompletes and type annotations. `pybind11-stubgen` is used to generate the stub files for the MEmilio Python Bindings and provide them as a separate stubs-only package.
+
+For installing stubs you first need to install our package `memilio.simulation` and the external dependency `pybind11-stubgen` for your python interpreter. Then run [generate_stubs.py](tools/generate_stubs.py) to generate the stubs-only package and install it as `memilio-stubs`, e.g. from the [current folder](.)
+
+```bash
+python ./tools/generate_stubs.py
+```
+
 ## Usage
 
 The package provides the following modules:

diff --git a/pycode/memilio-simulation/setup.py b/pycode/memilio-simulation/setup.py
@@ -20,8 +20,7 @@
     author_email='[email protected]', maintainer_email='[email protected]',
     url='https://github.com/SciCompMod/memilio',
     description='Part of MEmilio project, python bindings to the C++ libraries that contain the models and simulations.',
-    packages=find_packages(
-        where=os.path.dirname(os.path.abspath(__file__))),
+    packages=find_packages(where=os.path.dirname(os.path.abspath(__file__))),
     setup_requires=['cmake'],
     # need shared libs so there is one shared log level
     cmake_args=['-DMEMILIO_BUILD_SHARED_LIBS:BOOL=ON'],
@@ -36,4 +35,5 @@
             'numpy>=1.22,<1.25',
         ],
     },
-    long_description='', test_suite='memilio.simulation_test',)
+    long_description='', test_suite='memilio.simulation_test',
+)
diff --git a/pycode/memilio-simulation/tools/generate_stubs.py b/pycode/memilio-simulation/tools/generate_stubs.py
@@ -0,0 +1,72 @@
+import os
+import subprocess
+import sys
+import importlib.util
+import pkgutil
+import memilio.simulation
+
+setup_content = f"""
+from setuptools import setup, find_packages
+
+setup(
+    name='memilio-stubs',
+    version='0.1',
+    packages=['memilio-stubs'],
+    package_data={{
+        'memilio-stubs': ['simulation/*.pyi'],
+    }},
+)
+"""
+
+if __name__ == "__main__":
+
+    python_interpreter = sys.executable
+
+    # Check for needed packages. If it fails either pacakge is not installed or the wrong python interpreter is detected.
+    # For the latter try setting python_interpreter with full path
+    if importlib.util.find_spec('pybind11_stubgen') is None:
+        print('pybind11_stubgen is not installed')
+        exit()
+    if importlib.util.find_spec('memilio.simulation') is None:
+        print('memilio.simulation is not installed')
+        exit()
+
+    file_path = os.path.dirname(os.path.abspath(__file__))
+    package_dir = os.path.abspath(os.path.join(
+        file_path, "../../memilio-simulation-stubs"))
+    output_dir = os.path.join(package_dir, "memilio-stubs/simulation")
+    output_module_dir = os.path.join(output_dir, 'memilio')
+
+    # create folders, if they do not exist
+    try:
+        # output_module_dir is generated by pybind11_stubgen
+        os.makedirs(output_dir)
+    except:
+        pass
+
+    # get all model modules from memilio.simulation
+    # if package structure changes this needs to be adjusted
+    models = [m.name for m in pkgutil.iter_modules(
+        memilio.simulation.__path__)]
+
+    # generate stubs and moce them into correct folder with right name
+    # memilio-stubs/simulation module needs same structure as memilio/simulation
+    subprocess.check_call(
+        [python_interpreter, '-m', 'pybind11_stubgen', '--ignore-all-errors', '-o', output_dir, 'memilio._simulation'])
+    os.rename(os.path.join(output_module_dir, '_simulation.pyi'),
+              os.path.join(output_dir, '__init__.pyi'))
+
+    for model in models:
+        module_name = "memilio._simulation_" + model
+        subprocess.check_call(
+            [python_interpreter, '-m', 'pybind11_stubgen', '--ignore-all-errors', '-o', output_dir, module_name])
+        os.rename(os.path.join(output_module_dir, '_simulation_' + model + '.pyi'),
+                  os.path.join(output_dir, model + '.pyi'))
+
+    os.rmdir(output_module_dir)
+
+    # create setup.py and install package
+    with open(os.path.join(package_dir, "setup.py"), "w") as setup_file:
+        setup_file.write(setup_content)
+    subprocess.check_call(
+        [python_interpreter, '-m', 'pip', 'install', package_dir])