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.

[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