Implement mypy descriptor protocol for fields (es.Text and so on) #3251
Description
There is the usual magic going on when declaring documents with fields, but mypy isn't being told about it. For example, when defining a document as
import elasticsearch.dsl as es
class MyDoc(es.Document):
text = es.Text()
my_doc = MyDoc(text='foo')
reveal_type(my_doc.text) # Revealed type is "elasticsearch.dsl.field.Text"
reveal_type(MyDoc.text) # Revealed type is "elasticsearch.dsl.field.Text"
my_doc.text = 'bar' # error: Incompatible types in assignment (expression has type "str", variable has type "Text")
value: str = my_doc.text # error: Incompatible types in assignment (expression has type "Text", variable has type "str") Whereas the expected behavior would be that my_doc.text is typed as string, since it is a string as runtime (arguably str | None, but let me ignore that here).
Nowadays mypy fully supports type hints for descriptors, those can easily be used to have matching runtime types. Right now we can hack it in as
from __future__ import annotations
from typing import TYPE_CHECKING, overload
import elasticsearch.dsl as es
if not TYPE_CHECKING:
Text = es.Text
else:
class Text(es.Text):
@overload
def __get__(self, obj: None, cls: type[es.Document]) -> Text: ...
@overload
def __get__(self, obj: es.Document, cls: type[es.Document]) -> str: ...
def __get__(self, obj: es.Document | None, cls: type | None = None) -> Text | str:
raise RuntimeError('type stub only')
def __set__(self, obj: es.Document, value: str | Text) -> None:
pass
class MyDoc(es.Document):
text = Text()
my_doc = MyDoc(text='foo')
reveal_type(my_doc.text) # Revealed type is "builtins.str"
reveal_type(MyDoc.text) # Revealed type is "elasticsearch.dsl.field.Text"
my_doc.text = 'bar' # OK
value: str = my_doc.text # OKSo the proposal here would be to just add the __get__ and __set__ overrides to the fields in a TYPE_CHECKING block to give them the expected runtime type.