Best way to type fields for json_schema_extra?

[kwint]

We use the json_schema_extra to change the way our frontend visualizes input fields. I want to add typing to this so our developers know which options are available.

I expected this to work:

class CustomSchema(TypedDict):
    some_custom_option: str
    ui_option: list[Option]


class TestModel(BaseModel):
    val: str = Field(
        description="description",
        json_schema_extra=CustomSchema(
            some_custom_option="test value", ui_option=[Option.HIDDEN]
        ),
    )

But both mypy and pylance complain:

material_library/models/test.py:42: error: Argument "json_schema_extra" to "Field" has incompatible type "CustomSchema"; expected "JsonDict | Callable[[JsonDict], None] | None"  [arg-type]

I came up with this somewhat elegant solution, but feels more hacky then it should be?

from enum import StrEnum, auto

from pydantic import BaseModel, ConfigDict, Field
from pydantic.config import JsonDict


class Option(StrEnum):
    HIDDEN = auto()
    REFERENCE_ONLY = auto()


class CustomSchema(BaseModel):
    model_config = ConfigDict(use_enum_values=True)

    some_custom_option: str
    ui_option: list[Option]

    def __call__(self, schema: JsonDict) -> None:
        schema.update(self.model_dump())


class TestModel(BaseModel):
    val: str = Field(
        description="description",
        json_schema_extra=CustomSchema(
            some_custom_option="test value", ui_option=[Option.HIDDEN]
        ),
    )


print("schema: ", TestModel.model_json_schema())