refactor: create EmailEncoding enum

app/email_utils.py

@@ -1,5 +1,6 @@
 import base64
 import email
+import enum
 import os
 import quopri
 import random
@@ -673,7 +674,13 @@ def is_valid_email(email_address: str) -> bool:
     )
 
 
-def get_encoding(msg: Message) -> str:
+class EmailEncoding(enum.Enum):
+    BASE64 = "base64"
+    QUOTED = "quoted-printable"
+    NO = "no-encoding"
+
+
+def get_encoding(msg: Message) -> EmailEncoding:
     """
     Return the message encoding, possible values:
     - quoted-printable
@@ -682,21 +689,24 @@ def get_encoding(msg: Message) -> str:
     """
     cte = str(msg.get("content-transfer-encoding", "")).lower()
     if cte in ("", "7bit"):
-        return "7bit"
+        return EmailEncoding.NO
 
-    if cte in ("quoted-printable", "base64"):
-        return cte
+    if cte == "base64":
+        return EmailEncoding.BASE64
+
+    if cte == "quoted-printable":
+        return EmailEncoding.QUOTED
 
     LOG.exception("Unknown encoding %s", cte)
 
-    return "7bit"
+    return EmailEncoding.NO
 
 
-def encode_text(text: str, encoding: str = "7bit") -> str:
-    if encoding == "quoted-printable":
+def encode_text(text: str, encoding: EmailEncoding = EmailEncoding.NO) -> str:
+    if encoding == EmailEncoding.QUOTED:
         encoded = quopri.encodestring(text.encode("utf-8"))
         return str(encoded, "utf-8")
-    elif encoding == "base64":
+    elif encoding == EmailEncoding.BASE64:
         encoded = base64.b64encode(text.encode("utf-8"))
         return str(encoded, "utf-8")
     else:  # 7bit - no encoding

tests/test_email_utils.py

@@ -20,6 +20,7 @@ from app.email_utils import (
     normalize_reply_email,
     get_encoding,
     encode_text,
+    EmailEncoding,
 )
 from app.extensions import db
 from app.models import User, CustomDomain
@@ -422,24 +423,24 @@ def test_normalize_reply_email(flask_client):
 
 def test_get_encoding():
     msg = email.message_from_string("")
-    assert get_encoding(msg) == "7bit"
+    assert get_encoding(msg) == EmailEncoding.NO
 
     msg = email.message_from_string("Content-TRANSFER-encoding: Invalid")
-    assert get_encoding(msg) == "7bit"
+    assert get_encoding(msg) == EmailEncoding.NO
 
     msg = email.message_from_string("Content-TRANSFER-encoding: quoted-printable")
-    assert get_encoding(msg) == "quoted-printable"
+    assert get_encoding(msg) == EmailEncoding.QUOTED
 
     msg = email.message_from_string("Content-TRANSFER-encoding: base64")
-    assert get_encoding(msg) == "base64"
+    assert get_encoding(msg) == EmailEncoding.BASE64
 
 
 def test_encode_text():
     assert encode_text("") == ""
     assert encode_text("ascii") == "ascii"
-    assert encode_text("ascii", "base64") == "YXNjaWk="
-    assert encode_text("ascii", "quoted-printable") == "ascii"
+    assert encode_text("ascii", EmailEncoding.BASE64) == "YXNjaWk="
+    assert encode_text("ascii", EmailEncoding.QUOTED) == "ascii"
 
     assert encode_text("mèo méo") == "mèo méo"
-    assert encode_text("mèo méo", "base64") == "bcOobyBtw6lv"
-    assert encode_text("mèo méo", "quoted-printable") == "m=C3=A8o m=C3=A9o"
+    assert encode_text("mèo méo", EmailEncoding.BASE64) == "bcOobyBtw6lv"
+    assert encode_text("mèo méo", EmailEncoding.QUOTED) == "m=C3=A8o m=C3=A9o"