email.message.Message: Improve payload methods and Explain why __getitem__ isn't typed to return None

stdlib/email/message.pyi

@@ -3,17 +3,27 @@ from email import _ParamsType, _ParamType
 from email.charset import Charset
 from email.contentmanager import ContentManager
 from email.errors import MessageDefect
+from email.header import Header
 from email.policy import Policy
-from typing import Any, TypeVar, overload
-from typing_extensions import Self, TypeAlias
+from typing import Any, Protocol, TypeVar, overload
+from typing_extensions import Literal, Self, TypeAlias
 
 __all__ = ["Message", "EmailMessage"]
 
 _T = TypeVar("_T")
-
-_PayloadType: TypeAlias = list[Message] | str | bytes | bytearray
+_PayloadType: TypeAlias = Message | str
+_EncodedPayloadType: TypeAlias = Message | bytes
+_MultipartPayloadType: TypeAlias = list[_PayloadType]
 _CharsetType: TypeAlias = Charset | str | None
+# Type returned by Policy.header_fetch_parse, AnyOf[str | Header]
 _HeaderType: TypeAlias = Any
+_HeaderTypeParam: TypeAlias = str | Header
+
+class _SupportsEncodeToPayload(Protocol):
+    def encode(self, __encoding: str) -> _PayloadType | _MultipartPayloadType | _SupportsDecodeToPayload: ...
+
+class _SupportsDecodeToPayload(Protocol):
+    def decode(self, __encoding: str, __errors: str) -> _PayloadType | _MultipartPayloadType: ...
 
 class Message:
     policy: Policy  # undocumented
@@ -23,16 +33,43 @@ class Message:
     def is_multipart(self) -> bool: ...
     def set_unixfrom(self, unixfrom: str) -> None: ...
     def get_unixfrom(self) -> str | None: ...
-    def attach(self, payload: Message) -> None: ...
-    def get_payload(self, i: int | None = None, decode: bool = False) -> Any: ...  # returns _PayloadType | None
-    def set_payload(self, payload: _PayloadType, charset: _CharsetType = None) -> None: ...
+    def attach(self, payload: _PayloadType) -> None: ...
+    # `i: int` without a multipart payload results in an error
+    # `| Any`: can be None for cleared or unset payload, but annoying to check
+    @overload  # multipart
+    def get_payload(self, i: int, decode: Literal[True]) -> None: ...
+    @overload  # multipart
+    def get_payload(self, i: int, decode: Literal[False] = False) -> _PayloadType | Any: ...
+    @overload  # either
+    def get_payload(self, i: None = None, decode: Literal[False] = False) -> _PayloadType | _MultipartPayloadType | Any: ...
+    @overload  # not multipart
+    def get_payload(self, i: None = None, *, decode: Literal[True]) -> _EncodedPayloadType | Any: ...
+    @overload  # not multipart, IDEM but w/o kwarg
+    def get_payload(self, i: None, decode: Literal[True]) -> _EncodedPayloadType | Any: ...
+    # If `charset=None` and payload supports both `encode` AND `decode`, then an invalid payload could be passed, but this is unlikely
+    # Not[_SupportsEncodeToPayload]
+    @overload
+    def set_payload(
+        self, payload: _SupportsDecodeToPayload | _PayloadType | _MultipartPayloadType, charset: None = None
+    ) -> None: ...
+    @overload
+    def set_payload(
+        self,
+        payload: _SupportsEncodeToPayload | _SupportsDecodeToPayload | _PayloadType | _MultipartPayloadType,
+        charset: Charset | str,
+    ) -> None: ...
     def set_charset(self, charset: _CharsetType) -> None: ...
     def get_charset(self) -> _CharsetType: ...
     def __len__(self) -> int: ...
     def __contains__(self, name: str) -> bool: ...
     def __iter__(self) -> Iterator[str]: ...
+    # Same as `get` with `failobj=None`, but with the expectation that it won't return None in most scenarios
+    # This is important for protocols using __getitem__, like SupportsKeysAndGetItem
+    # Morally, the return type should be `AnyOf[_HeaderType, None]`,
+    # which we could spell as `_HeaderType | Any`,
+    # *but* `_HeaderType` itself is currently an alias to `Any`...
     def __getitem__(self, name: str) -> _HeaderType: ...
-    def __setitem__(self, name: str, val: _HeaderType) -> None: ...
+    def __setitem__(self, name: str, val: _HeaderTypeParam) -> None: ...
     def __delitem__(self, name: str) -> None: ...
     def keys(self) -> list[str]: ...
     def values(self) -> list[_HeaderType]: ...
@@ -46,7 +83,7 @@ class Message:
     @overload
     def get_all(self, name: str, failobj: _T) -> list[_HeaderType] | _T: ...
     def add_header(self, _name: str, _value: str, **_params: _ParamsType) -> None: ...
-    def replace_header(self, _name: str, _value: _HeaderType) -> None: ...
+    def replace_header(self, _name: str, _value: _HeaderTypeParam) -> None: ...
     def get_content_type(self) -> str: ...
     def get_content_maintype(self) -> str: ...
     def get_content_subtype(self) -> str: ...
@@ -100,7 +137,7 @@ class Message:
     ) -> None: ...
     def __init__(self, policy: Policy = ...) -> None: ...
     # The following two methods are undocumented, but a source code comment states that they are public API
-    def set_raw(self, name: str, value: _HeaderType) -> None: ...
+    def set_raw(self, name: str, value: _HeaderTypeParam) -> None: ...
     def raw_items(self) -> Iterator[tuple[str, _HeaderType]]: ...
 
 class MIMEPart(Message):