♻️ refactor(sqla): Model di by ModelParam

ProgramRipper · ProgramRipper · commit ad5f8d5bf26e · 2023-10-22T07:19:40.000Z
diff --git a/nonebot_plugin_orm/__init__.py b/nonebot_plugin_orm/__init__.py
@@ -222,6 +222,7 @@ def _init_logger():
 
 from .sql import *
 from .model import *
+from .param import *
 from .config import *
 from .migrate import *
 
diff --git a/nonebot_plugin_orm/model.py b/nonebot_plugin_orm/model.py
@@ -4,12 +4,10 @@
 from inspect import Parameter, Signature
 from typing import TYPE_CHECKING, Any, ClassVar
 
-from nonebot.params import Depends
+from sqlalchemy import Table, MetaData
 from nonebot import get_plugin_by_module_name
-from sqlalchemy import Table, MetaData, select
 from pydantic.typing import get_args, get_origin
 from sqlalchemy.orm import Mapped, DeclarativeBase
-from sqlalchemy.orm.decl_api import DeclarativeAttributeIntercept
 
 from .utils import DependsInner, get_annotations
 
@@ -30,47 +28,11 @@
 }
 
 
-class ModelMeta(DeclarativeAttributeIntercept):
-    if TYPE_CHECKING:
-        __signature__: Signature
-
-    def __new__(
-        mcs,
-        name: str,
-        bases: tuple[type, ...],
-        namespace: dict[str, Any],
-        **kwargs: Any,
-    ) -> ModelMeta:
-        from . import async_scoped_session
-
-        cls: ModelMeta = super().__new__(mcs, name, bases, namespace, **kwargs)
-
-        if not (signature := getattr(cls, "__signature__", None)):
-            return cls
-
-        async def dependency(
-            *, __session: async_scoped_session, **kwargs: Any
-        ) -> ModelMeta | None:
-            return await __session.scalar(select(cls).filter_by(**kwargs))
-
-        dependency.__signature__ = Signature(
-            (
-                Parameter(
-                    "_ModelMeta__session",
-                    Parameter.KEYWORD_ONLY,
-                    annotation=async_scoped_session,
-                ),
-                *signature.parameters.values(),
-            )
-        )
-
-        return Annotated[cls, Depends(dependency)]
-
-
-class Model(DeclarativeBase, metaclass=ModelMeta):
+class Model(DeclarativeBase):
     metadata = MetaData(naming_convention=NAMING_CONVENTION)
 
     if TYPE_CHECKING:
+        __signature__: Signature
         __table__: ClassVar[Table]
         __bind_key__: ClassVar[str]
 
@@ -90,8 +52,6 @@ def _setup_di(cls: type[Model]) -> None:
     """Get signature for NoneBot's dependency injection,
     and set annotations for SQLAlchemy declarative class.
     """
-    from . import async_scoped_session
-
     parameters: list[Parameter] = []
 
     annotations: dict[str, Any] = {}
diff --git a/nonebot_plugin_orm/param.py b/nonebot_plugin_orm/param.py
@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import sys
+from typing import Any, Union
+from inspect import Parameter, isclass
+from typing_extensions import Annotated
+from collections.abc import Iterator, Sequence, AsyncIterator
+
+from nonebot.matcher import Matcher
+from nonebot.dependencies import Param
+from pydantic.typing import get_args, get_origin
+from nonebot.params import DependParam, DefaultParam
+from sqlalchemy import Row, Result, ScalarResult, select
+from sqlalchemy.ext.asyncio import AsyncResult, AsyncScalarResult
+
+from .model import Model
+from .utils import Option, toclass, methodcall, compile_dependency
+
+if sys.version_info >= (3, 10):
+    from types import NoneType, UnionType
+else:
+    NoneType = type(None)
+    UnionType = None
+
+
+def parse_model_annotation(
+    anno: Any,
+) -> tuple[tuple[type[Model], ...], Option] | tuple[None, None]:
+    if isclass(anno) and issubclass(anno, Model):
+        return (anno,), Option(scalars=True, result=methodcall("one_or_none"))
+
+    origin, args = get_origin(anno), get_args(anno)
+
+    if not (origin and args):
+        return (None, None)
+
+    if origin is Annotated:
+        return parse_model_annotation(args[0])
+
+    if origin in (UnionType, Union) and len(args) == 2:
+        if args[0] is NoneType:
+            return parse_model_annotation(args[1])
+        elif args[1] is NoneType:
+            return parse_model_annotation(args[0])
+
+    if not isclass(origin):
+        return (None, None)
+
+    if origin is Row:
+        origin, args = tuple, get_args(args[0])
+
+    if origin is tuple and all(issubclass(arg, Model) for arg in map(toclass, args)):
+        return args, Option(result=methodcall("one_or_none"))
+
+    models, option = parse_model_annotation(args[0])
+    if not (models and option):
+        return (None, None)
+
+    if option.result == methodcall("all"):
+        if issubclass(Iterator, origin):
+            return models, Option(False, option.scalars, methodcall("partitions"))
+
+        if issubclass(AsyncIterator, origin):
+            return models, Option(True, option.scalars, methodcall("partitions"))
+
+    if option.result != methodcall("one_or_none"):
+        return (None, None)
+
+    if (
+        (not option.scalars and origin is Result)
+        or (option.scalars and origin is ScalarResult)
+        or issubclass(Iterator, origin)
+    ):
+        return models, Option(False, option.scalars)
+
+    if (
+        (not option.scalars and origin is AsyncResult)
+        or (option.scalars and origin is AsyncScalarResult)
+        or issubclass(AsyncIterator, origin)
+    ):
+        return models, Option(scalars=option.scalars)
+
+    if issubclass(Sequence, origin):
+        return models, Option(True, option.scalars, methodcall("all"))
+
+    return (None, None)
+
+
+class ModelParam(DependParam):
+    @classmethod
+    def _check_param(
+        cls, param: Parameter, allow_types: tuple[type[Param], ...]
+    ) -> Param | None:
+        models, option = parse_model_annotation(param.annotation)
+
+        if not (models and option):
+            return
+
+        stat = select(*models).where(
+            *(
+                getattr(model, name) == param.default
+                for model in models
+                for name, param in model.__signature__.parameters.items()
+            )
+        )
+
+        return super()._check_param(
+            param.replace(default=compile_dependency(stat, option)), allow_types
+        )
+
+    @classmethod
+    def _check_parameterless(
+        cls, value: Any, allow_types: tuple[type[Param], ...]
+    ) -> Param | None:
+        return
+
+
+Matcher.HANDLER_PARAM_TYPES = Matcher.HANDLER_PARAM_TYPES[:-1] + (
+    ModelParam,
+    DefaultParam,
+)
diff --git a/nonebot_plugin_orm/utils.py b/nonebot_plugin_orm/utils.py
@@ -5,16 +5,21 @@
 import logging
 from io import StringIO
 from pathlib import Path
-from typing import TypeVar
 from contextlib import suppress
 from functools import wraps, lru_cache
+from typing_extensions import Annotated
+from dataclasses import field, dataclass
+from typing import Any, TypeVar, NamedTuple
 from collections.abc import Callable, Iterable
+from inspect import Parameter, Signature, isclass
 from importlib.metadata import Distribution, PackageNotFoundError, distribution
 
 import click
 from nonebot.plugin import Plugin
 from nonebot.params import Depends
 from nonebot import logger, get_driver
+from pydantic.typing import get_args, get_origin
+from sqlalchemy.sql.selectable import ExecutableReturnsRows
 
 if sys.version_info >= (3, 9):
     from importlib.resources import files
@@ -83,6 +88,78 @@ def flush(self):
         pass
 
 
+@dataclass
+class Option:
+    stream: bool = True
+    scalars: bool = False
+    result: _methodcall | None = None
+    calls: list[_methodcall] = field(default_factory=list)
+
+
+class _methodcall(NamedTuple):
+    name: str
+    args: tuple
+    kwargs: dict[str, Any]
+
+
+def methodcall(name: str, args: tuple = (), kwargs: dict[str, Any] = {}) -> _methodcall:
+    return _methodcall(name, args, kwargs)
+
+
+def compile_dependency(statement: ExecutableReturnsRows, option: Option) -> Any:
+    from . import async_scoped_session
+
+    async def dependency(*, __session: async_scoped_session, **params: Any):
+        if option.stream:
+            result = await __session.stream(statement, params)
+        else:
+            result = await __session.execute(statement, params)
+
+        for call in option.calls:
+            result = getattr(result, call.name)(*call.args, **call.kwargs)
+
+        if option.scalars:
+            result = result.scalars()
+
+        if call := option.result:
+            result = getattr(result, call.name)(*call.args, **call.kwargs)
+
+            if option.stream:
+                result = await result
+
+        return result
+
+    dependency.__signature__ = Signature(
+        [
+            Parameter(
+                "__session", Parameter.KEYWORD_ONLY, annotation=async_scoped_session
+            ),
+            *(
+                Parameter(name, Parameter.KEYWORD_ONLY, default=depends)
+                for name, depends in statement.compile().params.items()
+                if isinstance(depends, DependsInner)
+            ),
+        ]
+    )
+
+    return Depends(dependency)
+
+
+def toclass(cls: Any) -> type:
+    if isclass(cls):
+        return cls
+
+    origin, args = get_origin(cls), get_args(cls)
+
+    if origin is None:
+        raise TypeError(f"{cls!r} is not a class or a generic type")
+
+    if origin is Annotated:
+        return toclass(args[0])
+
+    return origin
+
+
 def return_progressbar(func: Callable[_P, Iterable[_T]]) -> Callable[_P, Iterable[_T]]:
     log_level = get_driver().config.log_level
     if isinstance(log_level, str):
