Improve model extension import

Previously, it wasn't possible to import two model modules with the same name.
Now this is at least possible if they are in different locations.
Overwriting and importing a previously imported extension is still not supported.

python/examples/example_splines/ExampleSplines.ipynb

@@ -25,8 +25,6 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "import sys\n",
-    "from importlib import import_module\n",
     "from shutil import rmtree\n",
     "from tempfile import TemporaryDirectory\n",
     "from uuid import uuid1\n",
@@ -57,7 +55,7 @@
     "                parameters,\n",
     "                build_dir=build_dir,\n",
     "                model_name=model_name,\n",
-    "                **kwargs\n",
+    "                **kwargs,\n",
     "            )\n",
     "    else:\n",
     "        build_dir = os.path.join(BUILD_PATH, model_name)\n",
@@ -67,7 +65,7 @@
     "            parameters,\n",
     "            build_dir=build_dir,\n",
     "            model_name=model_name,\n",
-    "            **kwargs\n",
+    "            **kwargs,\n",
     "        )\n",
     "\n",
     "\n",
@@ -79,7 +77,7 @@
     "    model_name,\n",
     "    T=1,\n",
     "    discard_annotations=False,\n",
-    "    plot=True\n",
+    "    plot=True,\n",
     "):\n",
     "    if parameters is None:\n",
     "        parameters = {}\n",
@@ -89,13 +87,12 @@
     "    )\n",
     "    sbml_importer.sbml2amici(model_name, build_dir)\n",
     "    # Import the model module\n",
-    "    sys.path.insert(0, os.path.abspath(build_dir))\n",
-    "    model_module = import_module(model_name)\n",
+    "    model_module = amici.import_model_module(model_name, build_dir)\n",
     "    # Setup simulation timepoints and parameters\n",
     "    model = model_module.getModel()\n",
     "    for name, value in parameters.items():\n",
     "        model.setParameterByName(name, value)\n",
-    "    if isinstance(T, (int, float)):\n",
+    "    if isinstance(T, int | float):\n",
     "        T = np.linspace(0, T, 100)\n",
     "    model.setTimepoints([float(t) for t in T])\n",
     "    solver = model.getSolver()\n",
@@ -320,7 +317,7 @@
    ],
    "source": [
     "# Finally, we can simulate it in AMICI\n",
-    "model, rdata = simulate(sbml_model);"
+    "model, rdata = simulate(sbml_model)"
    ]
   },
   {

python/sdist/amici/__init__.py

@@ -7,14 +7,13 @@
 """
 
 import contextlib
-import datetime
+import importlib.util
 import importlib
 import os
 import re
 import sys
-import sysconfig
 from pathlib import Path
-from types import ModuleType as ModelModule
+from types import ModuleType
 from typing import Any
 from collections.abc import Callable
 
@@ -145,6 +144,8 @@ def get_model(self) -> amici.Model:
         def get_jax_model(self) -> JAXModel: ...
 
     AmiciModel = Union[amici.Model, amici.ModelPtr]
+else:
+    ModelModule = ModuleType
 
 
 class add_path:
@@ -182,6 +183,29 @@ def __exit__(self, exc_type, exc_value, traceback):
         sys.path = self.orginal_path
 
 
+def _module_from_path(module_name: str, module_path: Path | str) -> ModuleType:
+    """Import a module from a given path.
+
+    Import a module from a given path. The module is not added to
+    `sys.modules`. The `_self` attribute of the module is set to the module
+    itself.
+
+    :param module_name:
+        Name of the module.
+    :param module_path:
+        Path to the module file. Absolute or relative to the current working
+        directory.
+    """
+    module_path = Path(module_path).resolve()
+    if not module_path.is_file():
+        raise ModuleNotFoundError(f"Module file not found: {module_path}")
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    module._self = module
+    spec.loader.exec_module(module)
+    return module
+
+
 def import_model_module(
     module_name: str, module_path: Path | str
 ) -> ModelModule:
@@ -195,86 +219,29 @@ def import_model_module(
     :return:
         The model module
     """
-    module_path = str(module_path)
+    model_root = str(module_path)
 
     # ensure we will find the newly created module
     importlib.invalidate_caches()
 
     if not os.path.isdir(module_path):
-        raise ValueError(f"module_path '{module_path}' is not a directory.")
-
-    module_path = os.path.abspath(module_path)
-    ext_suffix = sysconfig.get_config_var("EXT_SUFFIX")
-    ext_mod_name = f"{module_name}._{module_name}"
-
-    # module already loaded?
-    if (m := sys.modules.get(ext_mod_name)) and m.__file__.endswith(
-        ext_suffix
-    ):
-        # this is the c++ extension we can't unload
-        loaded_file = Path(m.__file__)
-        needed_file = Path(
-            module_path,
-            module_name,
-            f"_{module_name}{ext_suffix}",
-        )
-        # if we import a matlab-generated model where the extension
-        #  is in a different directory
-        needed_file_matlab = Path(
-            module_path,
-            f"_{module_name}{ext_suffix}",
-        )
-        if not needed_file.exists():
-            if needed_file_matlab.exists():
-                needed_file = needed_file_matlab
-            else:
-                raise ModuleNotFoundError(
-                    f"Cannot find extension module for {module_name} in "
-                    f"{module_path}."
-                )
-
-        if not loaded_file.samefile(needed_file):
-            # this is not the right module, and we can't unload it
-            raise RuntimeError(
-                f"Cannot import extension for {module_name} from "
-                f"{module_path}, because an extension with the same name was "
-                f"has already been imported from {loaded_file.parent}. "
-                "Import the module with a different name or restart the "
-                "Python kernel."
-            )
-        # this is the right file, but did it change on disk?
-        t_imported = m._get_import_time()  # noqa: protected-access
-        t_modified = os.path.getmtime(m.__file__)
-        if t_imported < t_modified:
-            t_imp_str = datetime.datetime.fromtimestamp(t_imported).isoformat()
-            t_mod_str = datetime.datetime.fromtimestamp(t_modified).isoformat()
-            raise RuntimeError(
-                f"Cannot import extension for {module_name} from "
-                f"{module_path}, because an extension in the same location "
-                f"has already been imported, but the file was modified on "
-                f"disk. \nImported at {t_imp_str}\nModified at {t_mod_str}.\n"
-                "Import the module with a different name or restart the "
-                "Python kernel."
-            )
-
-    # unlike extension modules, Python modules can be unloaded
-    if module_name in sys.modules:
-        # if a module with that name is already in sys.modules, we remove it,
-        # along with all other modules from that package. otherwise, there
-        # will be trouble if two different models with the same name are to
-        # be imported.
-        del sys.modules[module_name]
-        # collect first, don't delete while iterating
-        to_unload = {
-            loaded_module_name
-            for loaded_module_name in sys.modules.keys()
-            if loaded_module_name.startswith(f"{module_name}.")
-        }
-        for m in to_unload:
-            del sys.modules[m]
-
-    with set_path(module_path):
-        return importlib.import_module(module_name)
+        raise ValueError(f"module_path '{model_root}' is not a directory.")
+
+    module_path = Path(model_root, module_name, "__init__.py")
+
+    # We may want to import a matlab-generated model where the extension
+    #  is in a different directory. This is not a regular use case. It's only
+    #  used in the amici tests and can be removed at any time.
+    #  The models (currently) use the default swig-import and require
+    #  modifying sys.path.
+    module_path_matlab = Path(model_root, f"{module_name}.py")
+    if not module_path.is_file() and module_path_matlab.is_file():
+        with set_path(model_root):
+            return _module_from_path(module_name, module_path_matlab)
+
+    module = _module_from_path(module_name, module_path)
+    module._self = module
+    return module
 
 
 class AmiciVersionError(RuntimeError):

python/sdist/amici/__init__.template.py

@@ -1,9 +1,12 @@
 """AMICI-generated module for model TPL_MODELNAME"""
 
+import datetime
+import os
 from pathlib import Path
 from typing import TYPE_CHECKING
 import amici
 
+
 if TYPE_CHECKING:
     from amici.jax import JAXModel
 
@@ -18,14 +21,39 @@
         "version currently installed."
     )
 
-from .TPL_MODELNAME import *  # noqa: F403, F401
-from .TPL_MODELNAME import getModel as get_model  # noqa: F401
+TPL_MODELNAME = amici._module_from_path(
+    "TPL_MODELNAME.TPL_MODELNAME", Path(__file__).parent / "TPL_MODELNAME.py"
+)
+for var in dir(TPL_MODELNAME):
+    if not var.startswith("_"):
+        globals()[var] = getattr(TPL_MODELNAME, var)
+get_model = TPL_MODELNAME.getModel
 
+# _self: this module; will be set during import
+TPL_MODELNAME._model_module = _self  # noqa: F821
 
-def get_jax_model() -> "JAXModel":
-    from .jax import JAXModel_TPL_MODELNAME
 
-    return JAXModel_TPL_MODELNAME()
+def get_jax_model() -> "JAXModel":
+    # If the model directory was meanwhile overwritten, this would load the
+    #  new version, which would not match the previously imported extension.
+    #  This is not allowed, as it would lead to inconsistencies.
+    jax_py_file = Path(__file__).parent / "jax.py"
+    jax_py_file = jax_py_file.resolve()
+    t_imported = TPL_MODELNAME._get_import_time()  # noqa: protected-access
+    t_modified = os.path.getmtime(jax_py_file)
+    if t_imported < t_modified:
+        t_imp_str = datetime.datetime.fromtimestamp(t_imported).isoformat()
+        t_mod_str = datetime.datetime.fromtimestamp(t_modified).isoformat()
+        raise RuntimeError(
+            f"Refusing to import {jax_py_file} which was changed since "
+            f"TPL_MODELNAME was imported. This is to avoid inconsistencies "
+            "between the different model implementations.\n"
+            f"Imported at {t_imp_str}\nModified at {t_mod_str}.\n"
+            "Import the module with a different name or restart the "
+            "Python kernel."
+        )
+    jax = amici._module_from_path("jax", jax_py_file)
+    return jax.JAXModel_TPL_MODELNAME()
 
 
 __version__ = "TPL_PACKAGE_VERSION"

python/tests/test_events.py

@@ -724,7 +724,7 @@ def test_handling_of_fixed_time_point_event_triggers():
     end
     """
     module_name = "test_events_time_based"
-    with TemporaryDirectory(prefix=module_name, delete=False) as outdir:
+    with TemporaryDirectory(prefix=module_name) as outdir:
         antimony2amici(
             ant_model,
             model_name=module_name,
@@ -765,7 +765,7 @@ def test_multiple_event_assignment_with_compartment():
     """
     # watch out for too long path names on windows ...
     module_name = "tst_mltple_ea_w_cmprtmnt"
-    with TemporaryDirectory(prefix=module_name, delete=False) as outdir:
+    with TemporaryDirectory(prefix=module_name) as outdir:
         antimony2amici(
             ant_model,
             model_name=module_name,