beanie [BUG] Incorrect type annotations for Document's `insert`, `save` and `replace`

Describe the bug when using Document.insert, save or replace vscode pylance report an error like this:

Argument missing for parameter "self" Pylance(reportGeneralTypeIssues)
(method) save: _Wrapped[..., Unknown, (self: Unknown, *args: Unknown, skip_actions: List[ActionDirections | str] | None = None, **kwargs: Unknown), Coroutine[Any, Any, Unknown]]

I found that it was due to the incorrect type annotation of the decorator, such as wrap_with_actions, save_state_after decorator.

If I remove its type comment, it's correct.

For example: https://github.com/roman-right/beanie/blob/main/beanie/odm/utils/state.py#L63C31

def save_state_after(f: Callable):  # remove the 'Callable' annotation
    @wraps(f)
    async def wrapper(self: "DocType", *args, **kwargs):
        result = await f(self, *args, **kwargs)
        self._save_state()
        return result

    return wrapper

This is probably the easiest way. If we want to use annotations, I guess we need to use ParamSpec and TypeVar for more detailed annotations. For example:

from functools import wraps
from typing import Callable, TypeVar
from typing_extensions import ParamSpec

P = ParamSpec("P")
R = TypeVar("R")


def example(some_arg):
    def decorator(f: Callable[P, R]) -> Callable[P, R]:
        @wraps(f)
        def wrapper(*args: P.args, **kwargs: P.kwargs):
            ...
            return f(*args, **kwargs)

        return wrapper

    return decorator

To Reproduce

import asyncio

from beanie import Document, init_beanie
from motor.motor_asyncio import AsyncIOMotorClient


class Foo(Document):
    name: str


async def main():
    client = AsyncIOMotorClient("mongodb://127.0.0.1:27017/")
    await init_beanie(
        database=client.testmango,
        document_models=[Foo],
    )

    foo = Foo(name="test")
    await foo.insert()  # Argument missing for parameter "self"
    await foo.save()  # Argument missing for parameter "self"
    await foo.replace()  # Argument missing for parameter "self"


if __name__ == "__main__":
    asyncio.run(main())

Expected behavior no error report

Additional context

Aug 28 '23 05:08 CMHopeSunshine

Was baffled by the error as well. will give it a shot

Sep 08 '23 15:09 mozTheFuzz

Hi, I'm also facing the same issue with pylance for Link[Document].

Whenever I try to assign a value to a Link[Document] with a new Document instance it triggers an error.

To reproduce:

from beanie import Document, Link

class Content(Document):
    id: UUID = Field(default_factory=uuid4)
    publisher = str
    category = str

class Book(Document):
    id: UUID = Field(default_factory=uuid4)
    book_pub: int
    book_content: Link[Content] | None


async def create_example_book() -> None:
      new_content = Content(
            publisher = "publisher_zebra",
            category = "science_fiction"
      )


      new_book_data = Book(
            book_pub=12345, 
            book_content=new_content, # error here (1) 
      )


(1) Error: Argument of type "Content" cannot be assigned to parameter "book_content" of type "Link[Content]" in function "__init__"  "Content" is incompatible with "Link[Content]"

Nov 13 '23 08:11 TyroneTang

I have the same issues

Nov 15 '23 07:11 helioascorreia

It seems to be due to the decorator conventions in this project. For example this decorator used in save

def validate_self_before(f: Callable):
    @wraps(f)
    async def wrapper(self: "DocType", *args, **kwargs):
        await self.validate_self(*args, **kwargs)
        return await f(self, *args, **kwargs)

    return wrapper

Because this wrapper is considered to be a module's function instead of a class method, self is considered to be a required positional argument, while we actually use it as a class method from an instance.

The way to make this work better with other linters is like this

def validate_self_before(f: Callable):
    @wraps(f)
    async def wrapper(*args, **kwargs):
        self = args[0]        
        await self.validate_self(args[1:], **kwargs)
        return await f(*args, **kwargs)
    return wrapper

I'd be more than happy to see the clean linter results from my IDEs too, but not sure if this is what the maintainers would like to do.

UPDATE: I guess this is something that the type checking/linting tools should update and reported it to pyright: https://github.com/microsoft/pyright/issues/6472

Nov 16 '23 08:11 jbkoh

However, shouldn't an instance calling a class method be considered to be the first argument always anyway?

@jbkoh tried following along with the explanation in microsoft/pyright#6472, but what did you mean in this question?

Just asking for clarification.

Dec 28 '23 07:12 ogtega

When you invoke instance.method(), instance is considered to be the first argument of method.

Dec 28 '23 07:12 jbkoh

To close the conversation, pyright folks said that we should annotate the arguments better as described here https://github.com/microsoft/pyright/issues/6472#issuecomment-1814022386

Dec 28 '23 07:12 jbkoh

I had a go at this and it does seem to work. I'm not sure it's 100% correct still, but at least it resolves the self error and still correctly identifies valid and invalid method arguments.

I changed: beanie/odm/actions.py like so:

from typing import ParamSpec, TypeVar, Concatenate

P = ParamSpec("P")
R = TypeVar("R")
S = TypeVar("S")

def wrap_with_actions(event_type: EventTypes):
    def decorator(f: Callable[Concatenate[S, P], R]) -> Callable[Concatenate[S, P], R]:
        @wraps(f)
        async def wrapper(
            self: S,
            *args: P.args,
            skip_actions: Optional[List[Union[ActionDirections, str]]] = None,
            **kwargs: P.kwargs,
        ):
            # the rest is the same

Then beanie/odm/utils/self.validation.py like so:

from functools import wraps
from typing import TYPE_CHECKING, Callable, ParamSpec, TypeVar, Concatenate

if TYPE_CHECKING:
    from beanie.odm.documents import DocType

P = ParamSpec("P")
R = TypeVar("R")

def validate_self_before(f: Callable[Concatenate["DocType", P], R]) -> Callable[Concatenate["DocType", P], R]:
    @wraps(f)
    async def wrapper(self: "DocType", *args: P.args, **kwargs: P.kwargs):
        await self.validate_self(*args, **kwargs)
        return await f(self, *args, **kwargs)

    return wrapper

and finally beanie/odm/utils/state.py like so:

from typing import ParamSpec, TypeVar, Concatenate

P = ParamSpec("P")
R = TypeVar("R")

def save_state_after(f: Callable[Concatenate["DocType", P], R]) -> Callable[Concatenate["DocType", P], R]:
    @wraps(f)
    async def wrapper(self: "DocType", *args: P.args, **kwargs: P.kwargs):
        result = await f(self, *args, **kwargs)
        self._save_state()
        return result

    return wrapper

Now I can do something like this, and there's no self error, it detects that link_rule is a valid param, and it detects does_not_exist is an invalid param:

I think it's still not perfect, because if I hover over link_rule it sees it as a function of type unknown:

(function) link_rule: Unknown

So it's not got the type info for the parameters of the method being wrapped. But it's still a lot better than it was before!

Anyone got any comments on this? Otherwise I might PR it as it is now

Feb 21 '24 13:02 TheBloke

Of course a big issue with the above is that Concatenate and ParamSpec require Python 3.10. So I guess any change like the above would need to be wrapped in a check that Python >= 3.10, and wouldn't help anyone using older versions.

Feb 21 '24 14:02 TheBloke

Of course a big issue with the above is that Concatenate and ParamSpec require Python 3.10. So I guess any change like the above would need to be wrapped in a check that Python >= 3.10, and wouldn't help anyone using older versions.

I'm tinkering with this a bit as well. Concatenate and ParamSpec are available in typing_extensions which is already a dependency.

I'm also trying to replace DocType with Self where applicable

UPD2

Nevermind, I did a stupid thing

Everything below is wrong. I found an error in my code and now all works as expected!

but it seems I found a bug in pyright because when type checking this

class Document(...):
    ...

    @wrap_with_actions(EventTypes.INSERT)
    @save_state_after
    @validate_self_before
    async def insert(
        self: Self,
        *,
        link_rule: WriteRules = WriteRules.DO_NOTHING,
        session: Optional[ClientSession] = None,
        skip_actions: Optional[List[Union[ActionDirections, str]]] = None,
    ) -> Self:
        ...

    async def create(
        self: Self,
        session: Optional[ClientSession] = None,
    ) -> Self:
        """
        The same as self.insert()
        :return: Document
        """
        reveal_type(self)
        reveal_type(self.insert)
        reveal_type(self.insert(session=session))
        reveal_type(await self.insert(session=session))
        return await self.insert(session=session)

I get this

  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:372:21 - information: Type of "self" is "Self@Document"
  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:373:21 - information: Type of "self.insert" is "(*, link_rule: WriteRules = WriteRules.DO_NOTHING, session: ClientSession | None = None, skip_actions: List[ActionDirections | str] | None = None) -> Coroutine[Any, Any, Document]"
  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:374:21 - information: Type of "self.insert(session=session)" is "Coroutine[Any, Any, Document]"
  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:375:21 - information: Type of "await self.insert(session=session)" is "Document"
  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:376:16 - error: Expression of type "Document" cannot be assigned to return type "Self@Document"
    Type "Document" cannot be assigned to type "Self@Document" (reportReturnType)

I don't see a reason why pyright thinks insert returns strict Document instead of Self@Document

UPD1:

beanie beanie copied to clipboard

[BUG] Incorrect type annotations for Document's `insert`, `save` and `replace`

beanie
beanie copied to clipboard