refactor: use asyncclick

2024-08-06 22:41:39 +08:00 · 2024-08-06 22:41:39 +08:00 · 4851ecfe82
commit 4851ecfe82
parent ad0e7006bc
4 changed files with 67 additions and 34 deletions
--- a/aerich/cli.py
+++ b/aerich/cli.py
@ -1,14 +1,11 @@
-import asyncio
 import os
-from functools import wraps
 from pathlib import Path
 from typing import Dict, List, cast

-import click
+import asyncclick as click
 import tomlkit
-from click import Context, UsageError
+from asyncclick import Context, UsageError
 from tomlkit.exceptions import NonExistentKey
-from tortoise import Tortoise

 from aerich import Command
 from aerich.enums import Color
@ -21,21 +18,6 @@ CONFIG_DEFAULT_VALUES = {
 }


-def coro(f):
-    @wraps(f)
-    def wrapper(*args, **kwargs) -> None:
-        loop = asyncio.get_event_loop()
-
-        # Close db connections at the end of all but the cli group function
-        try:
-            loop.run_until_complete(f(*args, **kwargs))
-        finally:
-            if f.__name__ not in ["cli", "init"]:
-                loop.run_until_complete(Tortoise.close_connections())
-
-    return wrapper
-
-
@click.group(context_settings={"help_option_names": ["-h", "--help"]})
@click.version_option(__version__, "-V", "--version")
@click.option(
@ -47,7 +29,6 @@ def coro(f):
 )
@click.option("--app", required=False, help="Tortoise-ORM app name.")
@click.pass_context
-@coro
 async def cli(ctx: Context, config, app) -> None:
    ctx.ensure_object(dict)
    ctx.obj["config_file"] = config
@ -83,7 +64,6 @@ async def cli(ctx: Context, config, app) -> None:
@click.option("--name", default="update", show_default=True, help="Migrate name.")
@click.option("--empty", default=False, is_flag=True, help="Generate empty migration file.")
@click.pass_context
-@coro
 async def migrate(ctx: Context, name) -> None:
    command = ctx.obj["command"]
    ret = await command.migrate(name)
@ -101,7 +81,6 @@ async def migrate(ctx: Context, name) -> None:
    help="Make migrations in transaction or not. Can be helpful for large migrations or creating concurrent indexes.",
 )
@click.pass_context
-@coro
 async def upgrade(ctx: Context, in_transaction: bool) -> None:
    command = ctx.obj["command"]
    migrated = await command.upgrade(run_in_transaction=in_transaction)
@ -133,7 +112,6 @@ async def upgrade(ctx: Context, in_transaction: bool) -> None:
@click.confirmation_option(
    prompt="Downgrade is dangerous, which maybe lose your data, are you sure?",
 )
-@coro
 async def downgrade(ctx: Context, version: int, delete: bool) -> None:
    command = ctx.obj["command"]
    try:
@ -146,7 +124,6 @@ async def downgrade(ctx: Context, version: int, delete: bool) -> None:

@cli.command(help="Show current available heads in migrate location.")
@click.pass_context
-@coro
 async def heads(ctx: Context) -> None:
    command = ctx.obj["command"]
    head_list = await command.heads()
@ -158,7 +135,6 @@ async def heads(ctx: Context) -> None:

@cli.command(help="List all migrate items.")
@click.pass_context
-@coro
 async def history(ctx: Context) -> None:
    command = ctx.obj["command"]
    versions = await command.history()
@ -189,7 +165,6 @@ async def history(ctx: Context) -> None:
    help="Folder of the source, relative to the project root.",
 )
@click.pass_context
-@coro
 async def init(ctx: Context, tortoise_orm, location, src_folder) -> None:
    config_file = ctx.obj["config_file"]

@ -233,7 +208,6 @@ async def init(ctx: Context, tortoise_orm, location, src_folder) -> None:
    show_default=True,
 )
@click.pass_context
-@coro
 async def init_db(ctx: Context, safe: bool) -> None:
    command = ctx.obj["command"]
    app = command.app
@ -257,7 +231,6 @@ async def init_db(ctx: Context, safe: bool) -> None:
    required=False,
 )
@click.pass_context
-@coro
 async def inspectdb(ctx: Context, table: List[str]) -> None:
    command = ctx.obj["command"]
    ret = await command.inspectdb(table)
--- a/aerich/utils.py
+++ b/aerich/utils.py
@ -6,7 +6,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Dict, Optional, Union

-from click import BadOptionUsage, ClickException, Context
+from asyncclick import BadOptionUsage, ClickException, Context
 from tortoise import BaseDBAsyncClient, Tortoise


--- a/poetry.lock
+++ b/poetry.lock
@ -28,6 +28,28 @@ files = [
 [package.dependencies]
 typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.9\""}

+[[package]]
+name = "anyio"
+version = "4.4.0"
+description = "High level compatibility layer for multiple asynchronous event loop implementations"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7"},
+    {file = "anyio-4.4.0.tar.gz", hash = "sha256:5aadc6a1bbb7cdb0bede386cac5e2940f5e2ff3aa20277e991cf028e0585ce94"},
+]
+
+[package.dependencies]
+exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""}
+idna = ">=2.8"
+sniffio = ">=1.1"
+typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""}
+
+[package.extras]
+doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
+test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
+trio = ["trio (>=0.23)"]
+
 [[package]]
 name = "async-timeout"
 version = "4.0.3"
@ -39,12 +61,27 @@ files = [
    {file = "async_timeout-4.0.3-py3-none-any.whl", hash = "sha256:7405140ff1230c310e51dc27b3145b9092d659ce68ff733fb0cefe3ee42be028"},
 ]

+[[package]]
+name = "asyncclick"
+version = "8.1.7.2"
+description = "Composable command line interface toolkit, async version"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "asyncclick-8.1.7.2-py3-none-any.whl", hash = "sha256:1ab940b04b22cb89b5b400725132b069d01b0c3472a9702c7a2c9d5d007ded02"},
+    {file = "asyncclick-8.1.7.2.tar.gz", hash = "sha256:219ea0f29ccdc1bb4ff43bcab7ce0769ac6d48a04f997b43ec6bee99a222daa0"},
+]
+
+[package.dependencies]
+anyio = "*"
+colorama = {version = "*", markers = "platform_system == \"Windows\""}
+
 [[package]]
 name = "asyncmy"
 version = "0.2.10rc1"
 description = "A fast asyncio MySQL driver"
 optional = true
-python-versions = "^3.7"
+python-versions = "<4.0,>=3.7"
 files = [
    {file = "asyncmy-0.2.10rc1.tar.gz", hash = "sha256:ba97b7f9b9719b6cb15169f0bffbf20be63767ff5052a24c3663a1d558bced5a"},
 ]
@ -364,6 +401,17 @@ files = [
 [package.extras]
 testing = ["hatch", "pre-commit", "pytest", "tox"]

+[[package]]
+name = "idna"
+version = "3.7"
+description = "Internationalized Domain Names in Applications (IDNA)"
+optional = false
+python-versions = ">=3.5"
+files = [
+    {file = "idna-3.7-py3-none-any.whl", hash = "sha256:82fee1fc78add43492d3a1898bfa6d8a904cc97d8427f683ed8e798d07761aa0"},
+    {file = "idna-3.7.tar.gz", hash = "sha256:028ff3aadf0609c1fd278d8ea3089299412a7a8b9bd005dd08b9f8285bcb5cfc"},
+]
+
 [[package]]
 name = "iniconfig"
 version = "2.0.0"
@ -816,6 +864,7 @@ files = [
    {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
    {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
    {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
+    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
    {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
    {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@ -895,6 +944,17 @@ files = [
    {file = "ruff-0.4.7.tar.gz", hash = "sha256:2331d2b051dc77a289a653fcc6a42cce357087c5975738157cd966590b18b5e1"},
 ]

+[[package]]
+name = "sniffio"
+version = "1.3.1"
+description = "Sniff out which async library your code is running under"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2"},
+    {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"},
+]
+
 [[package]]
 name = "stevedore"
 version = "5.2.0"
@ -936,7 +996,7 @@ name = "tortoise-orm"
 version = "0.21.3"
 description = "Easy async ORM for python, built with relations in mind"
 optional = false
-python-versions = ">=3.8,<4.0"
+python-versions = "<4.0,>=3.8"
 files = [
    {file = "tortoise_orm-0.21.3-py3-none-any.whl", hash = "sha256:9b8f8f8ba23a51f3407bfdc76cf9b2e5bc901ff07c7bec71250a83fa7724dab4"},
    {file = "tortoise_orm-0.21.3.tar.gz", hash = "sha256:d6e3a627915d4037d312f6ca0cb7d0bf6593630cf1da466df60c7c4c3128398e"},
@ -975,4 +1035,4 @@ asyncpg = ["asyncpg"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "1d8ab808c52f5859f9cae00b3d5127a5b91bbc6f8290faa1d6776f441252ca1b"
+content-hash = "ca5aa98cc3db69e15009ac541f161bab5bde8a53072a9dac39a78f58b5f8b183"
--- a/pyproject.toml
+++ b/pyproject.toml
@ -17,13 +17,13 @@ include = ["CHANGELOG.md", "LICENSE", "README.md"]
 [tool.poetry.dependencies]
 python = "^3.8"
 tortoise-orm = "*"
-click = "*"
 asyncpg = { version = "*", optional = true }
 asyncmy = { version = "^0.2.9", optional = true, allow-prereleases = true }
 pydantic = "^2.0"
 dictdiffer = "*"
 tomlkit = "*"

+asyncclick = "^8.1.7.2"
 [tool.poetry.group.dev.dependencies]
 ruff = "*"
 isort = "*"