utils: add types for safe_communicate

dmoj/utils/communicate.py

@@ -1,13 +1,16 @@
 import errno
 import os
 import select
+from typing import Dict, IO, List, Optional, Tuple
 
 from dmoj.error import OutputLimitExceeded
 
 _PIPE_BUF = getattr(select, 'PIPE_BUF', 512)
 
 
-def safe_communicate(proc, input=None, outlimit=None, errlimit=None):
+def safe_communicate(
+    proc, input: Optional[bytes] = None, outlimit: Optional[int] = None, errlimit: Optional[int] = None
+) -> Tuple[bytes, bytes]:
     if outlimit is None:
         outlimit = 10485760
     if errlimit is None:
@@ -19,12 +22,12 @@ def safe_communicate(proc, input=None, outlimit=None, errlimit=None):
         if not input:
             proc.stdin.close()
 
-    stdout = None  # Return
-    stderr = None  # Return
-    fd2file = {}
-    fd2output = {}
-    fd2length = {}
-    fd2limit = {}
+    stdout_list: Optional[List[bytes]] = None  # Return
+    stderr_list: Optional[List[bytes]] = None  # Return
+    fd2file: Dict[int, IO] = {}
+    fd2output: Dict[int, List[bytes]] = {}
+    fd2length: Dict[int, int] = {}
+    fd2limit: Dict[int, int] = {}
 
     poller = select.poll()
 
@@ -43,12 +46,12 @@ def close_unregister_and_remove(fd):
     select_POLLIN_POLLPRI = select.POLLIN | select.POLLPRI
     if proc.stdout:
         register_and_append(proc.stdout, select_POLLIN_POLLPRI)
-        fd2output[proc.stdout.fileno()] = stdout = []
+        fd2output[proc.stdout.fileno()] = stdout_list = []
         fd2length[proc.stdout.fileno()] = 0
         fd2limit[proc.stdout.fileno()] = outlimit
     if proc.stderr:
         register_and_append(proc.stderr, select_POLLIN_POLLPRI)
-        fd2output[proc.stderr.fileno()] = stderr = []
+        fd2output[proc.stderr.fileno()] = stderr_list = []
         fd2length[proc.stderr.fileno()] = 0
         fd2limit[proc.stderr.fileno()] = errlimit
 
@@ -63,6 +66,7 @@ def close_unregister_and_remove(fd):
 
         for fd, mode in ready:
             if mode & select.POLLOUT:
+                assert input
                 chunk = input[input_offset : input_offset + _PIPE_BUF]
                 try:
                     input_offset += os.write(fd, chunk)
@@ -90,8 +94,8 @@ def close_unregister_and_remove(fd):
                 close_unregister_and_remove(fd)
 
     # All data exchanged.  Translate lists into strings.
-    stdout = b''.join(stdout) if stdout is not None else b''
-    stderr = b''.join(stderr) if stderr is not None else b''
+    stdout = b''.join(stdout_list) if stdout_list is not None else b''
+    stderr = b''.join(stderr_list) if stderr_list is not None else b''
 
     proc.wait()
     return stdout, stderr