Annotate decorators that wrap Document methods (#679)

[bedlamzd]

Type checkers were complaining about missing self argument in decorated Document methods. This was caused by incomplete annotations of used decorators.

[bedlamzd]

I've seen tests/typing folder, but not sure what's the process for it. It seems the idea that mypy will error if there is anything wrong with typing in the test, right? Are they run as regular tests as well?

Let me know what kind of test you'd like to see and if simple usage of insert, save and replace is enough

[bedlamzd]

FYI typing_extensions has Self which is intended for this purpose. I can replace usages of DocType in Document class, if it isn't too out of scope. Or submit a new PR.

There is no difference between using two, though.

[bedlamzd]

Type annotation spec prohibits using named arguments along with typed **kwargs arguments.

See typing docs and PEP 612

One way to avoid using type: ignore is to write something like:

@wraps(f)
async def wrapper(
    self: "Document",
    *args: P.args,
    **kwargs: P.kwargs,
) -> R:
    skip_actions = cast(Optional[List[Union[ActionDirections, str]]], kwargs.get("skip_actions"))

    if skip_actions is None:
        skip_actions = []
    if skip_actions in kwargs:
        kwargs["skip_actions"] = skip_actions

    await ActionRegistry.run_actions(
        self,
        event_type=event_type,
        action_direction=ActionDirections.BEFORE,
        exclude=skip_actions,
    )

    result = await f(
        self,
        *args,
        **kwargs,
    )

    await ActionRegistry.run_actions(
        self,
        event_type=event_type,
        action_direction=ActionDirections.AFTER,
        exclude=skip_actions,
    )

    return result

But that doesn't look better in any way.

[roman-right]

Hi @bedlamzd,

Thank you very much for your work! Should this case be added to the mypy test suite, or was it a Pylance catch?

I use this test module for testing typing in usage: https://github.com/roman-right/beanie/tree/main/tests/typing.
This is where it runs: https://github.com/roman-right/beanie/blob/main/.pre-commit-config.yaml#L11-L18.

[bedlamzd]

@roman-right That's what I was wondering in the first comment

I don't remember how I've noticed it, likely an IDE check or pyright (I use it mostly). Mypy doesn't show errors because it infers Any type for Document.insert but I can write something like

from typing import (
    Awaitable,
    Protocol,
    assert_type,
)

from pymongo.client_session import ClientSession

from beanie import Document, WriteRules
from beanie.odm.actions import ActionDirections
from tests.typing.models import Test


class DocumentInsertSignature(Protocol):
    def __call__(
        self,
        *,
        link_rule: WriteRules = ...,
        session: ClientSession | None = ...,
        skip_actions: list[ActionDirections | str] | None = ...,
    ) -> Awaitable[Document]:
        ...


async def test_insert_signature() -> None:
    test = Test(
        foo="foo",
        bar="bar",
        baz="baz",
    )
    assert_type(test.insert, DocumentInsertSignature)

and mypy will do the check

You might've noticed that in the above signature returns Document and not Test as it should in reality. That is due to mypy weirdness, for some reason it can't infer subclass and fallbacks to Document. pyright doesn't have this issue and will raise an error here.

Also, protocol's __call__ is not async because in that case the return type is Coroutine[Any, Any, Document] which is incompatible with Awaitable in this context. But returning awaitable is basically the same thing as being async

[bedlamzd]

Turns out it's a lot simpler than I thought. This works as well and I think is better because works both for mypy and pyright

class DocumentClsInsertSignature(Protocol):
    def __call__(
        self,
        self_: DocType,
        /,
        *,
        link_rule: WriteRules = ...,
        session: ClientSession | None = ...,
        skip_actions: list[ActionDirections | str] | None = ...,
    ) -> Awaitable[DocType]:
        ...


def test_insert_signature() -> None:
    test_insert: DocumentClsInsertSignature = Test.insert  # noqa: F841

Suggested here

I will add tests a bit later

[bedlamzd]

@roman-right so while writing tests I found out that overloads don't work well with mypy, it always chooses async case. So I rewrote annotations and just ignore type errors in async wrappers.

You can see the changes in the fixup commit. Let me know when you've reviewed everything and I will squash it.

Overall I consider this ready, let me know if you want me to change anything.

[tristandeborde]

@bedlamzd Thanks a lot for the PR !
@roman-right do we have an estimate for when this PR could be reviewed and merged ?
Not trying to be pushy but this would be super helpful for my team - we're using Pylance :)
Since Pylance is VSCode's default Python linting option, surely it would greatly help with Beanie adoption !
Thanks in advance

[roman-right]

Hi @bedlamzd and @tristandeborde ,
Thank you very much for your work and following up. I plan to have this merged at this weekend

[rawnly]

hey any plan on releasing this?

[roman-right]

Hi everyone!
Merged. It will be released today