#20476: add a message_factory policy attribute to email.

4 files changed, 12 insertions, 11 deletions

diff --git a/Lib/email/_policybase.py b/Lib/email/_policybase.py
index c0d98a4..d699484 100644
--- a/Lib/email/_policybase.py
+++ b/Lib/email/_policybase.py
@@ -154,6 +154,8 @@ class Policy(_PolicyBase, metaclass=abc.ABCMeta):
                            them. This is used when the message is being
                            serialized by a generator. Default: True.
 
+    message_factory     -- the class to use to create new message objects.
+
     """
 
     raise_on_defect = False
@@ -161,6 +163,8 @@ class Policy(_PolicyBase, metaclass=abc.ABCMeta):
     cte_type = '8bit'
     max_line_length = 78
     mangle_from_ = False
+    # XXX To avoid circular imports, this is set in email.message.
+    message_factory = None
 
     def handle_defect(self, obj, defect):
         """Based on policy, either raise defect or call register_defect.
diff --git a/Lib/email/feedparser.py b/Lib/email/feedparser.py
index 2fa77d7..3d74978 100644
--- a/Lib/email/feedparser.py
+++ b/Lib/email/feedparser.py
@@ -24,7 +24,6 @@ __all__ = ['FeedParser', 'BytesFeedParser']
 import re
 
 from email import errors
-from email import message
 from email._policybase import compat32
 from collections import deque
 from io import StringIO
@@ -148,13 +147,7 @@ class FeedParser:
         self.policy = policy
         self._old_style_factory = False
         if _factory is None:
-            # What this should be:
-            #self._factory = policy.default_message_factory
-            # but, because we are post 3.4 feature freeze, fix with temp hack:
-            if self.policy is compat32:
-                self._factory = message.Message
-            else:
-                self._factory = message.EmailMessage
+            self._factory = policy.message_factory
         else:
             self._factory = _factory
             try:
diff --git a/Lib/email/message.py b/Lib/email/message.py
index c07da43..f4380d9 100644
--- a/Lib/email/message.py
+++ b/Lib/email/message.py
@@ -4,18 +4,17 @@
 
 """Basic message object for the email package object model."""
 
-__all__ = ['Message']
+__all__ = ['Message', 'EmailMessage']
 
 import re
 import uu
 import quopri
-import warnings
 from io import BytesIO, StringIO
 
 # Intrapackage imports
 from email import utils
 from email import errors
-from email._policybase import compat32
+from email._policybase import Policy, compat32
 from email import charset as _charset
 from email._encoded_words import decode_b
 Charset = _charset.Charset
@@ -1163,3 +1162,6 @@ class EmailMessage(MIMEPart):
         super().set_content(*args, **kw)
         if 'MIME-Version' not in self:
             self['MIME-Version'] = '1.0'
+
+# Set message_factory on Policy here to avoid a circular import.
+Policy.message_factory = Message
diff --git a/Lib/email/policy.py b/Lib/email/policy.py
index 35d0e69..5131311ac 100644
--- a/Lib/email/policy.py
+++ b/Lib/email/policy.py
@@ -7,6 +7,7 @@ from email._policybase import Policy, Compat32, compat32, _extend_docstrings
 from email.utils import _has_surrogates
 from email.headerregistry import HeaderRegistry as HeaderRegistry
 from email.contentmanager import raw_data_manager
+from email.message import EmailMessage
 
 __all__ = [
     'Compat32',
@@ -82,6 +83,7 @@ class EmailPolicy(Policy):
 
     """
 
+    message_factory = EmailMessage
     utf8 = False
     refold_source = 'long'
     header_factory = HeaderRegistry()