[BUG] Typing errors whenever `Link[SomeDoc]` is used (Pylance)

[Tom Jobbins]

Describe the bug
I am sure many people must have experienced this, but I can't find another issue for it (Except an old closed-as-stale one)

To Reproduce

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

class ModelToLink(Document):
    some_int: int

class ModelOne(Document):
    some_field: str
    link_field: Link[ModelToLink]
    
async def test_link():
    await init_beanie(database=client[DB_NAME], document_models=[ModelToLink, ModelOne])

    model_to_link = ModelToLink(some_int=1)

    model_one = ModelOne(some_field="test", link_field=model_to_link)
    # type error:
    # Argument of type "ModelToLink" cannot be assigned to parameter "link_field" of type "Link[ModelToLink]" in function "__init__"
    
    await model_one.save(link_rule=WriteRules.WRITE)
    
    assert model_one.link_field.some_int == 1
    # type error:
    # Cannot access member "some_int" for type "Link[ModelToLink]"
    #  Member "some_int" is unknown

Expected behavior
Links are followed to apply typing information as if they were the target of the link

Additional context

I started to see if I could fix this myself, and I found a way that half solves it.

In beanie/odm/fields.py:

T = TypeVar("T")

class Link(Generic[T]):
....
     if TYPE_CHECKING:
         def __get__(self, instance: Any, owner: Any) -> T:
            pass

This simple change half-solves the problem - now when in this example:

        assert model_one.link_field.some_int == 1

It will correctly detect some_int as type: int

But it does not fully work, because it will not detect when fields don't exist, eg:

        assert model_one.link_field.does_not_exist == 1

Will not raise any type error. does_not_exist will show as Unknown, but not type error is shown.

Can anyone help me get this fix fully working so that Link[SomeDoc] can evaluate all fields of SomeDoc with no typing errors, and also detect when fields do not exist?

Or am I missing something about how Link[..] works, and is there already some way to get this working?

I do love Link[], it's really powerful - but it's quite frustrating that it causes so many typing errors that have to have # type: ignore added, hiding real errors.

Thanks in advance

[Roman Right]

Hi! Thank you for the catch.

The main problem is that Pylance and mypy are different, and I paused Pylance support around 2 years ago when it started to implement new typing checks very frequently. It introduced new typing errors every week, stopping me from developing actual features. I think now it is more or less stable, so yes, I have to make Beanie support it again.

[Tom Jobbins]

Thank you! Yes I would say Pylance is pretty stable. There are releases no more than once a month, and nothing changed wrt to PyLance-Beanie in at least the last few releases (2023-12 -> 2024-02)

For now I worked around the issue by adding these helper functions:

T = TypeVar("T", bound=Document)

def as_link(instance: T) -> Link[T]:
    return cast(Link[T], instance)

def from_link(instance: Link[T]) -> T:
     return cast(T, instance)

So for example anywhere in my code I have a Link[X] field I can do

assert from_link(some_model.some_link_field).some_field == 1

It's not great though so would be awesome to get a proper fix in Beanie.

Thanks again.

[Roman Right]

Thank you for the context! I'll pick this up (the pylance support) after I finish Beanie and Bunnet sync. In a few weeks probably.

[Maxim]

The problem with Unknown is not because of the Link class but because LazyModel (which is parent of Document) defines __getattribute__. So type checker thinks "maybe this field is dynamic, so I just don't know".

Also in reality type of the linked field is Link[T] | T and which type it is exactly depends on link_rule parameter. So this technically is more correct:

class Link(Generic[DocType]):
    if TYPE_CHECKING:
        def __get__(self, instance, owner) -> Link[DocType] | DocType: ...

Otherwise linked field won't be type checked as Link.

Then whenever you use your models you'd have to disambiguate types:

from typing import reveal_type

class Nested(Document):
    foo: str


class MyModel(Document):
    linked: Link[Nested]

async def main():
    mm: MyModel = await MyModel.get("my_model_id", fetch_links=False)

    reveal_type(mm.linked)
    # information: Type of "mm.linked" is "Nested | Link[Nested]"

    reveal_type(mm.linked.foo)
    # information: Type of "mm.linked.foo" is "str | Unknown"
    # error: Cannot access member "foo" for type "Link[Nested]"

    if isinstance(mm.linked, Nested):
        reveal_type(mm.linked.foo)
        # information: Type of "mm.linked.foo" is "str"
        # NOTE: no error because of type narrowing via `isinstance` check

    reveal_type(mm.linked.fetch)
    # information: Type of "mm.linked.fetch" is "Any | Unknown | ((fetch_links: bool = False) -> Coroutine[Any, Any, Nested | Link[Nested]])"
    # NOTE: there is no error that `Nested` has no attribute `fetch` because it inherits `LazyModel`

I don't like this hack, because Link isn't a descriptor, but I don't know any other quick way either. Technically, the proper way would be to write a mypy plugin (as pydantic) that generates proper annotations, including __init__ etc. But I'd imagine it's too much just to make Link behave properly. Maybe if there is a way to hook into pydantic plugin, but I don't know anything about that.

Since @roman-right is maintainer of lazy-model, I decided to add this:

How to fix `LazyModel.__getattribute__`

Just hide it from type checker like that. Then undefined fields will raise errors as expected.

from typing import TYPE_CHECKING

class LazyModel(BaseModel):

    if TYPE_CHECKING:
        pass
    else:
        def __getattribute__(self, item):
            ...

I can open an issue in lazy-model, if needed.

[Maxim]

btw, duplicate #820

[tdeborde]

Thanks a lot @bedlamzd for the proposed fix! Did you end up opening an issue in Lazy Model?
Fixing this in Beanie would help my team delete some # type: ignore comments in our codebase :)

[Maxim]

@tristandeborde

Did you end up opening an issue in Lazy Model?

No, since I got no reply from the @roman-right. I assume he's busy with some other work, because there is almost no activity for the past few month.