From 65fcb89b4d774d02ccafea735737a106ba05f295 Mon Sep 17 00:00:00 2001
From: Daniel Kahn Gillmor <dkg@fifthhorseman.net>
Date: Fri, 1 Nov 2019 14:13:27 -0400
Subject: email-print-mime-structure: be typesafe

This adds enough typechecking that the following check passes:

    mypy --strict email-print-mimestructure

Signed-off-by: Daniel Kahn Gillmor <dkg@fifthhorseman.net>
---
 email-print-mime-structure | 54 +++++++++++++++++++++++++++-------------------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/email-print-mime-structure b/email-print-mime-structure
index 7adeb2b..185173f 100755
--- a/email-print-mime-structure
+++ b/email-print-mime-structure
@@ -29,44 +29,49 @@ Example:
 If you want to number the parts, i suggest piping the output through
 something like "cat -n"
 '''
-import email
 import sys
+import email
+import logging
 
-def print_part(z, prefix):
-    fname = '' if z.get_filename() is None else ' [' + z.get_filename() + ']'
-    cset = '' if z.get_charset() is None else ' (' + z.get_charset() + ')'
-    disp = z.get_params(None, header='Content-Disposition')
-    if (disp is None):
-        disposition = ''
-    else:
-        disposition = ''
+from typing import Optional, Union, List, Tuple, Any
+from email.charset import Charset
+from email.message import Message
+
+def print_part(z:Message, prefix:str) -> None:
+    ofname:Optional[str] = z.get_filename()
+    fname:str = '' if ofname is None else f' [{ofname}]'
+    ocharset:Union[Charset, str, None] = z.get_charset()
+    cset:str = '' if ocharset is None else f' ({ocharset})'
+    disp:Union[List[Tuple[str,str]], List[str], None] = z.get_params(None, header='Content-Disposition')
+    disposition:str = ''
+    if (disp is not None):
         for d in disp:
             if d[0] in [ 'attachment', 'inline' ]:
                 disposition = ' ' + d[0]
+    nbytes:int
     if z.is_multipart():
+        # FIXME: it looks like we are counting chars here, not bytes:
         nbytes = len(z.as_string())
     else:
-        nbytes = len(z.get_payload())
+        payload:Union[List[Message], str, bytes, None] = z.get_payload()
+        if not isinstance(payload, (str,bytes)):
+            raise TypeError(f'expected payload to be either str or bytes, got {type(payload)}')
+        nbytes = len(payload)
 
-    print('{}{}{}{}{} {:d} bytes'.format(
-        prefix,
-        z.get_content_type(),
-        cset,
-        disposition,
-        fname,
-        nbytes,
-    ))
+    print(f'{prefix}{z.get_content_type()}{cset}{disposition}{fname} {nbytes} bytes')
 
-def test(z, prefix=''):
+def test(z:Message, prefix:str='') -> None:
     if (z.is_multipart()):
         print_part(z, prefix+'┬╴')
         if prefix.endswith('└'):
             prefix = prefix.rpartition('└')[0] + ' '
         if prefix.endswith('├'):
             prefix = prefix.rpartition('├')[0] + '│'
-        parts = z.get_payload()
+        parts:Union[List[Message], str, bytes, None] = z.get_payload()
+        if not isinstance(parts, list):
+            raise TypeError(f'parts was {type(parts)}, expected List[Message]')
         i = 0
-        while (i < parts.__len__()-1):
+        while (i < len(parts)-1):
             test(parts[i], prefix + '├')
             i += 1
         test(parts[i], prefix + '└')
@@ -74,4 +79,9 @@ def test(z, prefix=''):
     else:
         print_part(z, prefix+'─╴')
 
-test(email.message_from_file(sys.stdin), '└')
+msg:Union[Message, str, int, Any] = email.message_from_file(sys.stdin)
+
+if isinstance(msg, Message):
+    test(msg, '└')
+else:
+    logging.error('Input was not an e-mail message')
-- 
cgit v1.2.3