Make email.message.Message generic over the header type (#11732)

Co-authored-by: Avasam <samuel.06@hotmail.com>

stdlib/email/message.pyi

@@ -3,22 +3,25 @@ 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, Literal, Protocol, TypeVar, overload
+from typing import Any, Generic, Literal, Protocol, TypeVar, overload
 from typing_extensions import Self, TypeAlias
 
 __all__ = ["Message", "EmailMessage"]
 
 _T = TypeVar("_T")
+# Type returned by Policy.header_fetch_parse, often str or Header.
+_HeaderT = TypeVar("_HeaderT", default=str)
+_HeaderParamT = TypeVar("_HeaderParamT", default=str)
+# Represents headers constructed by HeaderRegistry. Those are sub-classes
+# of BaseHeader and another header type.
+_HeaderRegistryT = TypeVar("_HeaderRegistryT", default=Any)
+_HeaderRegistryParamT = TypeVar("_HeaderRegistryParamT", default=Any)
+
 _PayloadType: TypeAlias = Message | str
 _EncodedPayloadType: TypeAlias = Message | bytes
 _MultipartPayloadType: TypeAlias = list[_PayloadType]
 _CharsetType: TypeAlias = Charset | str | None
-# Type returned by Policy.header_fetch_parse, often str or Header.
-_HeaderType: TypeAlias = Any
-# Type accepted by Policy.header_store_parse.
-_HeaderTypeParam: TypeAlias = str | Header | Any
 
 class _SupportsEncodeToPayload(Protocol):
     def encode(self, encoding: str, /) -> _PayloadType | _MultipartPayloadType | _SupportsDecodeToPayload: ...
@@ -26,10 +29,7 @@ class _SupportsEncodeToPayload(Protocol):
 class _SupportsDecodeToPayload(Protocol):
     def decode(self, encoding: str, errors: str, /) -> _PayloadType | _MultipartPayloadType: ...
 
-# TODO: This class should be generic over the header policy and/or the header
-# value types allowed by the policy. This depends on PEP 696 support
-# (https://github.com/python/typeshed/issues/11422).
-class Message:
+class Message(Generic[_HeaderT, _HeaderParamT]):
     policy: Policy  # undocumented
     preamble: str | None
     epilogue: str | None
@@ -70,24 +70,23 @@ class Message:
     # 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: _HeaderTypeParam) -> None: ...
+    # so using "the Any trick" instead.
+    def __getitem__(self, name: str) -> _HeaderT | Any: ...
+    def __setitem__(self, name: str, val: _HeaderParamT) -> None: ...
     def __delitem__(self, name: str) -> None: ...
     def keys(self) -> list[str]: ...
-    def values(self) -> list[_HeaderType]: ...
-    def items(self) -> list[tuple[str, _HeaderType]]: ...
+    def values(self) -> list[_HeaderT]: ...
+    def items(self) -> list[tuple[str, _HeaderT]]: ...
     @overload
-    def get(self, name: str, failobj: None = None) -> _HeaderType | None: ...
+    def get(self, name: str, failobj: None = None) -> _HeaderT | None: ...
     @overload
-    def get(self, name: str, failobj: _T) -> _HeaderType | _T: ...
+    def get(self, name: str, failobj: _T) -> _HeaderT | _T: ...
     @overload
-    def get_all(self, name: str, failobj: None = None) -> list[_HeaderType] | None: ...
+    def get_all(self, name: str, failobj: None = None) -> list[_HeaderT] | None: ...
     @overload
-    def get_all(self, name: str, failobj: _T) -> list[_HeaderType] | _T: ...
+    def get_all(self, name: str, failobj: _T) -> list[_HeaderT] | _T: ...
     def add_header(self, _name: str, _value: str, **_params: _ParamsType) -> None: ...
-    def replace_header(self, _name: str, _value: _HeaderTypeParam) -> None: ...
+    def replace_header(self, _name: str, _value: _HeaderParamT) -> None: ...
     def get_content_type(self) -> str: ...
     def get_content_maintype(self) -> str: ...
     def get_content_subtype(self) -> str: ...
@@ -141,14 +140,14 @@ 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: _HeaderTypeParam) -> None: ...
-    def raw_items(self) -> Iterator[tuple[str, _HeaderType]]: ...
+    def set_raw(self, name: str, value: _HeaderParamT) -> None: ...
+    def raw_items(self) -> Iterator[tuple[str, _HeaderT]]: ...
 
-class MIMEPart(Message):
+class MIMEPart(Message[_HeaderRegistryT, _HeaderRegistryParamT]):
     def __init__(self, policy: Policy | None = None) -> None: ...
-    def get_body(self, preferencelist: Sequence[str] = ("related", "html", "plain")) -> Message | None: ...
-    def iter_attachments(self) -> Iterator[Message]: ...
-    def iter_parts(self) -> Iterator[Message]: ...
+    def get_body(self, preferencelist: Sequence[str] = ("related", "html", "plain")) -> MIMEPart[_HeaderRegistryT] | None: ...
+    def iter_attachments(self) -> Iterator[MIMEPart[_HeaderRegistryT]]: ...
+    def iter_parts(self) -> Iterator[MIMEPart[_HeaderRegistryT]]: ...
     def get_content(self, *args: Any, content_manager: ContentManager | None = None, **kw: Any) -> Any: ...
     def set_content(self, *args: Any, content_manager: ContentManager | None = None, **kw: Any) -> None: ...
     def make_related(self, boundary: str | None = None) -> None: ...