refactor: use asyncclick to replace click (#6)

* refactor: use asyncclick to replace click

* refactor: move logic of get tortoise orm config to utils

* Remove unused pass

Author: waketzheng

CHANGELOG.md

@@ -4,10 +4,12 @@
 
 ### 0.2.0
 
+- Use `asyncclick` instead of `click` ([#6])
 - Only install tomlkit for Python version less than 3.11 ([#5])
 - Migrate lint tool from isort+black to ruff ([#5])
 - Drop support for Python3.8 ([#4])
 
+[#6]: https://github.com/tortoise/tortoise-cli/pull/6
 [#5]: https://github.com/tortoise/tortoise-cli/pull/5
 [#4]: https://github.com/tortoise/tortoise-cli/pull/4
 

poetry.lock

@@ -10,8 +10,8 @@ include = ["CHANGELOG.md", "LICENSE", "README.md"]
 requires-python = ">=3.9"
 dependencies = [
     "tortoise-orm",
-    "click",
     "ptpython",
+    "asyncclick (>=8.1.7,<9.0.0)",
     "tomlkit (>=0.11.4,<1.0.0); python_version < '3.11'",
 ]
 

pyproject.toml

@@ -1,35 +1,23 @@
 from __future__ import annotations
 
-import asyncio
 import contextlib
-import os
 import sys
-from functools import wraps
-from pathlib import Path
+from collections.abc import AsyncGenerator
 
-import click
+import asyncclick as click
 from ptpython.repl import embed
-from tortoise import Tortoise
+from tortoise import Tortoise, connections
 
 from tortoise_cli import __version__, utils
 
-if sys.version_info >= (3, 11):
-    import tomllib
-else:
-    import tomlkit as tomllib
 
-
-def coro(f):
-    @wraps(f)
-    def wrapper(*args, **kwargs):
-        loop = asyncio.get_event_loop()
-        try:
-            loop.run_until_complete(f(*args, **kwargs))
-        finally:
-            if f.__name__ != "cli":
-                loop.run_until_complete(Tortoise.close_connections())
-
-    return wrapper
+@contextlib.asynccontextmanager
+async def aclose_tortoise() -> AsyncGenerator[None]:
+    try:
+        yield
+    finally:
+        if Tortoise._inited:
+            await connections.close_all()
 
 
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
@@ -40,16 +28,8 @@ def wrapper(*args, **kwargs):
     help="TortoiseORM config dictionary path, like settings.TORTOISE_ORM",
 )
 @click.pass_context
-@coro
-async def cli(ctx: click.Context, config: str | None):
-    if (
-        not config
-        and not (config := os.getenv("TORTOISE_ORM"))
-        and (p := Path("pyproject.toml")).exists()
-    ):
-        doc = tomllib.loads(p.read_text("utf-8"))
-        config = doc.get("tool", {}).get("aerich", {}).get("tortoise_orm", "")
-    if not config:
+async def cli(ctx: click.Context, config: str | None) -> None:
+    if not config and not (config := utils.tortoise_orm_config()):
         raise click.UsageError(
             "You must specify TORTOISE_ORM in option or env, or config file pyproject.toml from config of aerich",
             ctx=ctx,
@@ -61,19 +41,19 @@ async def cli(ctx: click.Context, config: str | None):
 
 @cli.command(help="Start a interactive shell.")
 @click.pass_context
-@coro
-async def shell(ctx: click.Context):
-    with contextlib.suppress(EOFError, ValueError):
-        await embed(
-            globals=globals(),
-            title="Tortoise Shell",
-            vi_mode=True,
-            return_asyncio_coroutine=True,
-            patch_stdout=True,
-        )
-
-
-def main():
+async def shell(ctx: click.Context) -> None:
+    async with aclose_tortoise():
+        with contextlib.suppress(EOFError, ValueError):
+            await embed(
+                globals=globals(),
+                title="Tortoise Shell",
+                vi_mode=True,
+                return_asyncio_coroutine=True,
+                patch_stdout=True,
+            )
+
+
+def main() -> None:
     if sys.path[0] != ".":
         sys.path.insert(0, ".")
     cli()

tortoise_cli/cli.py

@@ -1,6 +1,27 @@
 import importlib
+import os
+import sys
+from pathlib import Path
 
-from click import BadOptionUsage, ClickException, Context
+from asyncclick import BadOptionUsage, ClickException, Context
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomlkit as tomllib
+
+
+def tortoise_orm_config(file="pyproject.toml") -> str:
+    """
+    get tortoise orm config from os environment variable or aerich item in pyproject.toml
+
+    :param file: toml file that aerich item loads from it
+    :return: module path and var name that store the tortoise config, e.g.: 'settings.TORTOISE_ORM'
+    """
+    if not (config := os.getenv("TORTOISE_ORM", "")) and (p := Path(file)).exists():
+        doc = tomllib.loads(p.read_text("utf-8"))
+        config = doc.get("tool", {}).get("aerich", {}).get("tortoise_orm", "")
+    return config
 
 
 def get_tortoise_config(ctx: Context, config: str) -> dict: