Repository: tortoise/aerich
Pull request #458 (2025-05-29 23:52:45)

fix: refactor module importing to work with byte compiled migrations
In some cases an application may be installed with byte compiled migrations, to handle this scenario correctly implement module based import/search logic using pkgutil.

jameshilliard: @waketzheng Does this look good to merge?

Patch:
diff --git a/aerich/__init__.py b/aerich/__init__.py
index bd9bd3e..349eadb 100644
--- a/aerich/__init__.py
+++ b/aerich/__init__.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import os
+import pkgutil
 import platform
 from collections.abc import Generator
 from contextlib import AbstractAsyncContextManager
@@ -20,10 +20,13 @@ from aerich.inspectdb.sqlite import InspectSQLite
 from aerich.migrate import MIGRATE_TEMPLATE, Migrate
 from aerich.models import Aerich
 from aerich.utils import (
+    file_module_info,
     get_app_connection,
     get_app_connection_name,
     get_models_describe,
     import_py_file,
+    import_py_module,
+    py_module_path,
 )
 
 if TYPE_CHECKING:
@@ -179,10 +182,16 @@ class Command(AbstractAsyncContextManager):
         await self.close()
 
     async def _upgrade(
-        self, conn: BaseDBAsyncClient, version_file: str, fake: bool = False
+        self,
+        conn: BaseDBAsyncClient,
+        version_file: str,
+        fake: bool = False,
+        version_module: pkgutil.ModuleInfo | None = None,
     ) -> None:
-        file_path = Path(Migrate.migrate_location, version_file)
-        m = import_py_file(file_path)
+        if version_module is not None:
+            m = import_py_module(version_module)
+        else:
+            m = import_py_file(Path(Migrate.migrate_location, version_file))
         upgrade = m.upgrade
         if not fake:
             await conn.execute_script(await upgrade(conn))
@@ -194,7 +203,8 @@ class Command(AbstractAsyncContextManager):
 
     async def upgrade(self, run_in_transaction: bool = True, fake: bool = False) -> list[str]:
         migrated = []
-        for version_file in Migrate.get_all_version_files():
+        for version_module in Migrate.get_all_version_modules():
+            version_file = version_module.name + ".py"
             try:
                 exists = await Aerich.exists(version=version_file, app=self.app)
             except OperationalError:
@@ -203,10 +213,10 @@ class Command(AbstractAsyncContextManager):
                 app_conn_name = get_app_connection_name(self.tortoise_config, self.app)
                 if run_in_transaction:
                     async with in_transaction(app_conn_name) as conn:
-                        await self._upgrade(conn, version_file, fake=fake)
+                        await self._upgrade(conn, version_file, fake, version_module)
                 else:
                     app_conn = get_app_connection(self.tortoise_config, self.app)
-                    await self._upgrade(app_conn, version_file, fake=fake)
+                    await self._upgrade(app_conn, version_file, fake, version_module)
                 migrated.append(version_file)
         return migrated
 
@@ -229,8 +239,8 @@ class Command(AbstractAsyncContextManager):
             async with in_transaction(
                 get_app_connection_name(self.tortoise_config, self.app)
             ) as conn:
-                file_path = Path(Migrate.migrate_location, file)
-                m = import_py_file(file_path)
+                module_info = file_module_info(Migrate.migrate_location, Path(file).stem)
+                m = import_py_module(module_info)
                 downgrade = m.downgrade
                 downgrade_sql = await downgrade(conn)
                 if not downgrade_sql.strip():
@@ -239,7 +249,7 @@ class Command(AbstractAsyncContextManager):
                     await conn.execute_script(downgrade_sql)
                 await version_obj.delete()
                 if delete:
-                    os.unlink(file_path)
+                    py_module_path(module_info).unlink()
                 ret.append(file)
         return ret
 
diff --git a/aerich/migrate.py b/aerich/migrate.py
index 192bc4a..dd0e04c 100644
--- a/aerich/migrate.py
+++ b/aerich/migrate.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import contextlib
 import importlib
-import os
+import pkgutil
 import re
 from collections.abc import Iterable
 from datetime import datetime
@@ -24,8 +24,9 @@ from aerich.utils import (
     get_app_connection,
     get_dict_diff_by_key,
     get_models_describe,
-    import_py_file,
+    import_py_module,
     is_default_function,
+    py_module_path,
     run_async,
 )
 
@@ -66,20 +67,22 @@ class Migrate:
         return next(filter(lambda x: x.get("name") == name, fields))
 
     @classmethod
-    def get_all_version_files(cls) -> list[str]:
-        def get_file_version(file_name: str) -> str:
-            return file_name.split("_")[0]
+    def get_all_version_modules(cls) -> list[pkgutil.ModuleInfo]:
+        def get_file_version(module: pkgutil.ModuleInfo) -> str:
+            return module.name.split("_")[0]
 
-        def is_version_file(file_name: str) -> bool:
-            if not file_name.endswith("py"):
-                return False
-            if "_" not in file_name:
+        def is_version_file(module: pkgutil.ModuleInfo) -> bool:
+            if "_" not in module.name:
                 return False
-            return get_file_version(file_name).isdigit()
+            return get_file_version(module).isdigit()
 
-        files = filter(is_version_file, os.listdir(cls.migrate_location))
+        files = filter(is_version_file, pkgutil.iter_modules([str(cls.migrate_location)]))
         return sorted(files, key=lambda x: int(get_file_version(x)))
 
+    @classmethod
+    def get_all_version_files(cls) -> list[str]:
+        return [m.name + ".py" for m in cls.get_all_version_modules()]
+
     @classmethod
     def _get_model(cls, model: str) -> type[Model]:
         return Tortoise.apps[cls.app].get(model)  # type: ignore
@@ -141,17 +144,17 @@ class Migrate:
     async def _generate_diff_py(cls, name, no_input: bool = False) -> str | None:
         content = cls._get_diff_file_content()
         version = await cls.generate_version(name)  # '<num>_<date>_<name>.py'
-        conflict_files = [
-            version_file
-            for version_file in cls.get_all_version_files()
-            if version_file.startswith(version.split("_")[0])
+        conflict_modules = [
+            version_module
+            for version_module in cls.get_all_version_modules()
+            if version_module.name.startswith(version.split("_")[0])
         ]
-        if conflict_files:
-            if len(conflict_files) == 1:
-                file = Path(cls.migrate_location, conflict_files[0])
+        if conflict_modules:
+            if len(conflict_modules) == 1:
+                file = py_module_path(conflict_modules[0])
                 tip = f"Migration file exists({file}). Do you want to remove it?"
             else:
-                tip = f"Migration file exists({conflict_files}). Do you want to remove them?"
+                tip = f"Migration file exists({[py_module_path(m) for m in conflict_modules]}). Do you want to remove them?"
             overwrite = no_input or click.prompt(
                 tip,
                 default=False,
@@ -161,8 +164,8 @@ class Migrate:
             if not overwrite:
                 return None
             # delete same version files
-            for version_file in conflict_files:
-                os.unlink(Path(cls.migrate_location, version_file))
+            for version_module in conflict_modules:
+                py_module_path(version_module).unlink()
 
         Path(cls.migrate_location, version).write_text(content, encoding="utf-8")
         return version
@@ -428,12 +431,12 @@ class Migrate:
             return False
         # For postgresql, if a unique_together was created when generating the table, it is
         # a constraint. And if it was created after table generated, it will be a unique index.
-        migrate_files = cls.get_all_version_files()
+        migrate_files = cls.get_all_version_modules()
         if len(migrate_files) < 2:
             return True
         pattern = re.compile(rf' "?{index_name}"? ')
         for filename in reversed(migrate_files[1:]):
-            module = import_py_file(Path(cls.migrate_location, filename))
+            module = import_py_module(filename)
             upgrade_sql = run_async(module.upgrade, None)
             if pattern.search(upgrade_sql):
                 line = [i for i in upgrade_sql.splitlines() if pattern.search(i)][0]
diff --git a/aerich/utils.py b/aerich/utils.py
index 9c670e3..40d9c84 100644
--- a/aerich/utils.py
+++ b/aerich/utils.py
@@ -2,9 +2,11 @@ from __future__ import annotations
 
 import importlib.util
 import os
+import pkgutil
 import re
 import sys
 from collections.abc import Awaitable, Callable, Generator
+from importlib.machinery import FileFinder
 from pathlib import Path
 from types import ModuleType
 from typing import Any, TypeVar
@@ -116,6 +118,13 @@ def is_default_function(string: Any) -> re.Match | None:
     return re.match(r"^<function.+>$", str(string or ""))
 
 
+def file_module_info(path: str | Path, name: str) -> pkgutil.ModuleInfo:
+    for module_info in pkgutil.iter_modules([str(path)]):
+        if module_info.name == name:
+            return module_info
+    raise FileNotFoundError(f"Module {name} not found in path {path}")
+
+
 def import_py_file(file: str | Path) -> ModuleType:
     module_name, file_ext = os.path.splitext(os.path.split(file)[-1])
     spec = importlib.util.spec_from_file_location(module_name, file)
@@ -124,6 +133,30 @@ def import_py_file(file: str | Path) -> ModuleType:
     return module
 
 
+def import_py_module(module_info: pkgutil.ModuleInfo) -> ModuleType:
+    module_finder: FileFinder
+    name: str
+    ispkg: bool
+    module_finder, name, ispkg = module_info  # type:ignore[assignment]
+    module_finder.invalidate_caches()
+    spec = module_finder.find_spec(name)
+    module = importlib.util.module_from_spec(spec)  # type:ignore[arg-type]
+    spec.loader.exec_module(module)  # type:ignore[union-attr]
+    return module
+
+
+def py_module_path(module_info: pkgutil.ModuleInfo) -> Path:
+    module_finder: FileFinder
+    name: str
+    ispkg: bool
+    module_finder, name, ispkg = module_info  # type:ignore[assignment]
+    module_finder.invalidate_caches()
+    spec = module_finder.find_spec(name)
+    if not spec or not spec.origin or not Path(spec.origin).is_file():
+        raise FileNotFoundError(f"Module {name} not found in {module_finder.path}.")
+    return Path(spec.origin)
+
+
 def get_dict_diff_by_key(
     old_fields: list[dict], new_fields: list[dict], key="through"
 ) -> Generator[tuple]:


Test patch:
diff --git a/tests/test_migrate.py b/tests/test_migrate.py
index 9df0ffd..10f5f79 100644
--- a/tests/test_migrate.py
+++ b/tests/test_migrate.py
@@ -1222,7 +1222,7 @@ def test_sort_files_containing_non_migrations(mocker):
         return_value=[
             "1_datetime_update.py",
             "11_datetime_update.py",
-            "10_datetime_update.py",
+            "10_datetime_update.pyc",
             "2_datetime_update.py",
             "not_a_migration.py",
             "999.py",