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

[CMHopeSunshine]

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

[mozTheFuzz]

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

[TyroneTang]

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]"

[helioascorreia]

I have the same issues

[jbkoh]

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: microsoft/pyright#6472

[ogtega]

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.

[jbkoh]

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

[jbkoh]

To close the conversation, pyright folks said that we should annotate the arguments better as described here
microsoft/pyright#6472 (comment)

[TheBloke]

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

[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.

[bedlamzd]

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:

More on `Self` type erasure

        reveal_type(type(self))
        reveal_type(type(self).insert)

outputs

  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:371:21 - information: Type of "type(self)" is "type[Self@Document]"
  /Users/bedlamzd/Repositories/beanie/beanie/odm/documents.py:372:21 - information: Type of "type(self).insert" is "(Document, *, link_rule: WriteRules = WriteRules.DO_NOTHING, session: ClientSession | None = None, skip_actions: List[ActionDirections | str] | None = None) -> Coroutine[Any, Any, Document]"

So when accessing methods pyright for some reason resolves the class type

[TheBloke]

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

Oh yes, great point! I forgot about that.

So are you already working on a PR for this? I won't do one if you already have one you're working on.

[bedlamzd]

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

Oh yes, great point! I forgot about that.

So are you already working on a PR for this? I won't do one if you already have one you're working on.

I don't have a PR for this, so you can go ahead if you have something ready. I just stumbled upon this yesterday and thought it'd be interesting to look around.

But I can start working on a PR, I think I'm almost done with the decorators. The only thing left is to annotate async cases properly

[bedlamzd]

FYI I've opened a PR for this issue #886

Sorry it took few days, got distracted