Fix type hints for UserMessage

src/magentic/chat_model/message.py

@@ -31,6 +31,9 @@
 
 from magentic.function_call import FunctionCall
 
+if TYPE_CHECKING:
+    from magentic.typing import NonStringSequence
+
 PlaceholderT = TypeVar("PlaceholderT", covariant=True)
 
 
@@ -162,16 +165,22 @@ def format(self, **kwargs: Any) -> Self:
         return self
 
 
-UserMessageContentBlock: TypeAlias = str | ImageBytes | ImageUrl
-UserMessageContentBlockT = TypeVar(
-    "UserMessageContentBlockT", bound=UserMessageContentBlock, covariant=True
-)
+UserMessageContentBlock: TypeAlias = ImageBytes | ImageUrl
 UserMessageContentT = TypeVar(
     "UserMessageContentT",
     bound=str
-    | Sequence[UserMessageContentBlock | Placeholder[UserMessageContentBlock]],
+    | Sequence[str | UserMessageContentBlock | Placeholder[UserMessageContentBlock]],
     covariant=True,
 )
+UserMessageContentBlockT = TypeVar(
+    "UserMessageContentBlockT", bound=UserMessageContentBlock
+)
+UserMessageContentBlockT2 = TypeVar(
+    "UserMessageContentBlockT2", bound=UserMessageContentBlock
+)
+# These allow type hinting that a `str` _might_ be part of the union
+StrT = TypeVar("StrT", str, str)
+StrT2 = TypeVar("StrT2", str, str)
 
 
 class UserMessage(Message[UserMessageContentT], Generic[UserMessageContentT]):
@@ -187,20 +196,15 @@ def format(self: "UserMessage[str]", **kwargs: Any) -> "UserMessage[str]": ...
 
     @overload
     def format(
-        self: "UserMessage[Sequence[UserMessageContentBlockT]]", **kwargs: Any
-    ) -> "UserMessage[Sequence[UserMessageContentBlockT]]": ...
-
-    @overload
-    def format(
-        self: "UserMessage[Sequence[Placeholder[UserMessageContentBlockT]]]",
+        self: "UserMessage[StrT | NonStringSequence[StrT2 | UserMessageContentBlockT | Placeholder[UserMessageContentBlockT2]]]",
         **kwargs: Any,
-    ) -> "UserMessage[Sequence[UserMessageContentBlockT]]": ...
+    ) -> "UserMessage[StrT | Sequence[StrT2 | UserMessageContentBlockT | UserMessageContentBlockT2]]": ...
 
     def format(
-        self: "UserMessage[str | Sequence[UserMessageContentBlockT | Placeholder[UserMessageContentBlockT]]]",
+        self: "UserMessage[StrT | NonStringSequence[StrT2 | UserMessageContentBlockT | Placeholder[UserMessageContentBlockT2]]]",
         **kwargs: Any,
-    ) -> "UserMessage[str | Sequence[UserMessageContentBlockT]]":
-        if isinstance(self.content, str | Placeholder):
+    ) -> "UserMessage[StrT | Sequence[StrT2 | UserMessageContentBlockT | UserMessageContentBlockT2]]":
+        if isinstance(self.content, str):
             return UserMessage(self.content.format(**kwargs))
         if isinstance(self.content, Iterable):
             return UserMessage([block.format(**kwargs) for block in self.content])  # type: ignore[misc]

src/magentic/typing.py

@@ -1,7 +1,26 @@
 import inspect
 import types
 from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, TypeGuard, TypeVar, Union, get_args, get_origin
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Protocol,
+    TypeGuard,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
+
+T_co = TypeVar("T_co", covariant=True)
+
+if TYPE_CHECKING:
+    # Cannot be defined at runtime because Protocol cannot inherit from non-Protocol
+    class NonStringSequence(Sequence[T_co], Protocol[T_co]):  # type: ignore[misc]
+        """Protocol that matches Sequences except for `str`."""
+
+        # HACK: Works because `__contains__` method of `str` does not match `Sequence`
+        # See: https://github.com/python/typing/issues/256#issuecomment-1442633430
 
 
 def is_union_type(type_: type) -> bool:

tests/chat_model/test_message.py

@@ -11,6 +11,7 @@
     AssistantMessage,
     FunctionResultMessage,
     ImageBytes,
+    ImageUrl,
     Placeholder,
     SystemMessage,
     ToolResultMessage,
@@ -19,6 +20,9 @@
 )
 from magentic.function_call import FunctionCall
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 def test_placeholder_format():
     class Country(BaseModel):
@@ -111,6 +115,82 @@ def test_user_message_format():
     assert user_message_formatted == UserMessage("Hello world")
 
 
+def test_user_message_format_type_hints():
+    if TYPE_CHECKING:  # Avoid runtime error for None missing `format` method
+        assert_type(cast(UserMessage[Literal["x"]], None).format(), UserMessage[str])
+        # mypy does not convert `Literal` to `str` in these cases but pyright does
+        # assert_type(cast(UserMessage[Sequence[Literal["x"]]], None).format(), UserMessage[Sequence[str]])  # noqa: ERA001
+        # assert_type(cast(UserMessage[Sequence[Literal["x"] | ImageBytes]], None).format(), UserMessage[Sequence[str | ImageBytes]])  # noqa: ERA001
+
+        assert_type(cast(UserMessage[str], None).format(), UserMessage[str])
+
+        assert_type(
+            cast(UserMessage[Sequence[str]], None).format(), UserMessage[Sequence[str]]
+        )
+        assert_type(
+            cast(UserMessage[str | Sequence[str]], None).format(),
+            UserMessage[str | Sequence[str]],
+        )
+
+        assert_type(
+            cast(UserMessage[Sequence[ImageBytes]], None).format(),
+            UserMessage[Sequence[ImageBytes]],
+        )
+        assert_type(
+            cast(UserMessage[Sequence[str | ImageBytes]], None).format(),
+            UserMessage[Sequence[str | ImageBytes]],
+        )
+        assert_type(
+            cast(UserMessage[str | Sequence[str | ImageBytes]], None).format(),
+            UserMessage[str | Sequence[str | ImageBytes]],
+        )
+
+        assert_type(
+            cast(UserMessage[Sequence[Placeholder[ImageUrl]]], None).format(),
+            UserMessage[Sequence[ImageUrl]],
+        )
+        assert_type(
+            cast(UserMessage[Sequence[str | Placeholder[ImageUrl]]], None).format(),
+            UserMessage[Sequence[str | ImageUrl]],
+        )
+        assert_type(
+            cast(
+                UserMessage[Sequence[ImageBytes | Placeholder[ImageUrl]]], None
+            ).format(),
+            UserMessage[Sequence[ImageBytes | ImageUrl]],
+        )
+        assert_type(
+            cast(
+                UserMessage[Sequence[str | ImageBytes | Placeholder[ImageUrl]]], None
+            ).format(),
+            UserMessage[Sequence[str | ImageBytes | ImageUrl]],
+        )
+
+        assert_type(
+            cast(UserMessage[str | Sequence[Placeholder[ImageUrl]]], None).format(),
+            UserMessage[str | Sequence[ImageUrl]],
+        )
+        assert_type(
+            cast(
+                UserMessage[str | Sequence[str | Placeholder[ImageUrl]]], None
+            ).format(),
+            UserMessage[str | Sequence[str | ImageUrl]],
+        )
+        assert_type(
+            cast(
+                UserMessage[str | Sequence[ImageBytes | Placeholder[ImageUrl]]], None
+            ).format(),
+            UserMessage[str | Sequence[ImageBytes | ImageUrl]],
+        )
+        assert_type(
+            cast(
+                UserMessage[str | Sequence[str | ImageBytes | Placeholder[ImageUrl]]],
+                None,
+            ).format(),
+            UserMessage[str | Sequence[str | ImageBytes | ImageUrl]],
+        )
+
+
 def test_assistant_message_usage():
     assistant_message = AssistantMessage("Hello")
     assert assistant_message.usage is None
@@ -140,6 +220,8 @@ class Country(BaseModel):
     assert_type(assistant_message_formatted.content, Country)
     assert assistant_message_formatted == AssistantMessage(Country(name="USA"))
 
+
+def test_assistant_message_format_type_hints():
     if TYPE_CHECKING:  # Avoid runtime error for None missing `format` method
         assert_type(
             cast(AssistantMessage[Literal["x"]], None).format(), AssistantMessage[str]