♻️ refactor: split init_orm()

Author: ProgramRipper

nonebot_plugin_orm/__init__.py

@@ -84,14 +84,8 @@
 
 
 @_driver.on_startup
-async def init_orm():
-    global _session_factory
-
-    _init_engines()
-    _init_table()
-    _session_factory = sa_asyncio.async_sessionmaker(
-        _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
-    )
+async def init_orm() -> None:
+    _init_orm()
 
     with migrate.AlembicConfig(stdout=StreamToLogger()) as alembic_config:
         if plugin_config.alembic_startup_check:
@@ -102,6 +96,16 @@ async def init_orm():
             await greenlet_spawn(migrate.stamp, alembic_config)
 
 
+def _init_orm():
+    global _session_factory
+
+    _init_engines()
+    _init_table()
+    _session_factory = sa_asyncio.async_sessionmaker(
+        _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
+    )
+
+
 @wraps(lambda: None)  # NOTE: for dependency injection
 def get_session(**local_kw: Any) -> sa_asyncio.AsyncSession:
     try:

nonebot_plugin_orm/__main__.py

@@ -7,6 +7,7 @@
 
 import click
 from alembic.script import Script
+from sqlalchemy.util import greenlet_spawn
 
 from . import migrate
 from .config import plugin_config
@@ -205,7 +206,7 @@ def stamp(*args, **kwargs) -> None:
 @orm.command()
 @click.argument("rev", default="current")
 @click.pass_obj
-def edit(*args, **kwargs):
+def edit(*args, **kwargs) -> None:
     """使用 $EDITOR 编辑修订文件。"""
     return migrate.edit(*args, **kwargs)
 
@@ -218,12 +219,14 @@ def ensure_version(*args, **kwargs) -> None:
     return migrate.ensure_version(*args, **kwargs)
 
 
-def main():
-    from . import _init_table, _init_engines
+def main(*args, **kwargs) -> None:
+    from . import _init_orm
 
-    _init_engines()
-    _init_table()
-    orm(prog_name="nb orm")
+    if not (args or kwargs):
+        kwargs["prog_name"] = "nb orm"
+
+    _init_orm()
+    orm(*args, **kwargs)
 
 
 if __name__ == "__main__":