"json_schema_extra" and Pylance/Pyright issues

Question

I need to add metadata to fields of a pydantic model in a way that I could change the metadata. I ended up with the following solution:

class Foo(BaseModel):
    a: str = Field(
        meta_field=("some extra data a"),  # pyright: ignore
    )
    b: str = Field(
        meta_field=("some extra data b"),  # pyright: ignore
    )
    c: str = Field(
        meta_field=("some extra data c"),  # pyright: ignore
    )

    @classmethod
    def summarize_meta_fields(cls, **kwargs) -> dict[str, str]:
        schema = cls.model_json_schema()
        return {
            k: schema["properties"][k]["meta_field"] for k in schema["properties"].keys()
        }


def configure_meta_data(**kwargs) -> None:
    for k in kwargs:
        if k not in Foo.model_fields:
            raise ValueError(f"Field {k} not found in SummeryTube model")
        Foo.model_fields[k].json_schema_extra["meta_field"] = kwargs[k]

My problem is that in VScode, I get the following error:

with the following text:

Object of type "(JsonDict) -> None" is not subscriptablePylancereportIndexIssue
Object of type "None" is not subscriptablePylancereportOptionalSubscript

How can I refactor the code and mitigate this warning? Or should I simply ignore it as the code behaves as I expect.

Answer 1

The problem is that pyright thinks json_schema_extra may have type Callable[[JsonDict], None] | None. And it's... correct. As per documentation:

You can pass a dict to json_schema_extra to add extra information to the JSON schema:

[...]

You can pass a Callable to json_schema_extra to modify the JSON schema with a function

So the static type of that attribute supports a dict or a callable indeed, and you can only use bracket access with the dict part.

Here's the definition in source:

from typing import Callable, Dict, List, TypeAlias

JsonValue: TypeAlias = Union[int, float, str, bool, None, List['JsonValue'], 'JsonDict']
JsonDict: TypeAlias = Dict[str, JsonValue]

class FieldInfo(_repr.Representation):
    ...
    json_schema_extra: JsonDict | Callable[[JsonDict], None] | None
    ...

(FieldInfo is the type of Foo.model_fields[k], as explained in the docstring)

So pyright is correct here, you either need to handle all cases or assert the expected type.

def configure_meta_data(**kwargs) -> None:
    for k in kwargs:
        if k not in Foo.model_fields:
            raise ValueError(f"Field {k} not found in SummeryTube model")
        field_info = Foo.model_fields[k]
        if isinstance(field_info.json_schema_extra, dict):
            field_info.json_schema_extra["meta_field"] = kwargs[k]
        else:
            ...  # Do something else - error, warn, ignore, ...

You may also want to handle other cases (perhaps set to an empty dict if None and wrap a callable with a decorator according to your desired changes).

Sorry, I'm too lazy to install pydantic and pyright locally, and playground doesn't have pydantic installed. This code should work but wasn't tested.

NB (by the OP): The tested, and correct, solution is:

def configure_meta_data(**kwargs) -> None:
    for k in kwargs:
        if k not in Foo.model_fields:
            raise ValueError(f"Field {k} not found in SummeryTube model")
        json_extra_data = Foo.model_fields[k].json_schema_extra
        if not isinstance(json_extra_data, dict):
            raise ValueError(f"Field {k} has no json_schema_extra")
        json_extra_data["meta_field"] = kwargs[k]