T = TypeVar("T")
U = TypeVar("U")

class Key(Generic[T]):
    ...

class CustomDict(dict[Key[Any] | str, Any]):
    @overload  # type: ignore[override]
    def get(self, __key: Key[T]) -> T | None:
        ...

    @overload
    def get(self, __key: Key[T], __default: U) -> T | U:
        ...

    @overload
    def get(self, __key: str) -> Any | None:
        ...

    @overload
    def get(self, __key: str, __default: Any) -> Any:
        ...

    def get(self, __key: Key[Any] | str, __default: Any = None) -> Any:
        """Forward to super implementation."""
        return super().get(__key, __default)

    # overloads for __getitem__, setdefault, pop
    # ...

    @overload  # type: ignore[override]
    def __setitem__(self, key: Key[T], value: T) -> None:
        ...

    @overload
    def __setitem__(self, key: str, value: Any) -> None:
        ...

    def __setitem__(self, key: Key[Any] | str, value: Any) -> None:
        """Forward to super implementation."""
        return super().__setitem__(key, value)

d = CustomDict()
key = Key[int]()
other_key = "other"
assert_type(d.get(key), int | None)
assert_type(d.get("other"), Any | None)

def a2(d: CustomDict) -> None:
    if (var := d.get("arg")) is None:
        var = d["arg"] = {}
        reveal_type(var)