commands: add mypy types

dmoj/cli.py

@@ -2,8 +2,9 @@
 import readline
 import shlex
 import sys
-from typing import cast
+from typing import List, cast
 
+from dmoj.commands.base_command import GradedSubmission
 from dmoj.error import InvalidCommandException
 from dmoj.judge import Judge
 from dmoj.packet import PacketManager
@@ -61,6 +62,8 @@ def close(self):
 
 
 class LocalJudge(Judge):
+    graded_submissions: List[GradedSubmission]
+
     def __init__(self):
         super().__init__(cast(PacketManager, LocalPacketManager(self)))
         self.submission_id_counter = 0

dmoj/commands/base_command.py

@@ -4,19 +4,24 @@
 import sys
 import tempfile
 from collections import OrderedDict
-from typing import Dict
+from typing import Dict, NoReturn, Optional, TYPE_CHECKING, Tuple
 
 from dmoj.error import InvalidCommandException
 from dmoj.executors import executors
 from dmoj.utils.ansi import print_ansi
 
+if TYPE_CHECKING:
+    from dmoj.cli import LocalJudge
+
+GradedSubmission = Tuple[str, str, str, float, int]
+
 
 class CommandArgumentParser(argparse.ArgumentParser):
-    def error(self, message):
+    def error(self, message: str) -> NoReturn:
         self.print_usage(sys.stderr)
         raise InvalidCommandException
 
-    def exit(self, status=0, message=None):
+    def exit(self, status: int = 0, message: Optional[str] = None) -> NoReturn:
         if message:
             self._print_message(message, sys.stderr)
         raise InvalidCommandException
@@ -26,62 +31,62 @@ class Command:
     name = 'command'
     help = ''
 
-    def __init__(self, judge):
+    def __init__(self, judge: 'LocalJudge') -> None:
         self.judge = judge
         self.arg_parser = CommandArgumentParser(prog=self.name, description=self.help)
         self._populate_parser()
 
-    def get_source(self, source_file):
+    def get_source(self, source_file: str) -> str:
         try:
             with open(os.path.realpath(source_file)) as f:
                 return f.read()
         except Exception as io:
             raise InvalidCommandException(str(io))
 
-    def get_submission_data(self, submission_id):
+    def get_submission_data(self, submission_id: int) -> GradedSubmission:
         # don't wrap around
         if submission_id > 0:
             try:
                 return self.judge.graded_submissions[submission_id - 1]
             except IndexError:
                 pass
 
-        raise InvalidCommandException("invalid submission '%d'" % submission_id)
+        raise InvalidCommandException(f"invalid submission '{submission_id}'")
 
-    def open_editor(self, lang, src=b''):
+    def open_editor(self, lang: str, src: str = '') -> str:
         file_suffix = '.' + executors[lang].Executor.ext
         editor = os.environ.get('EDITOR')
         if editor:
-            with tempfile.NamedTemporaryFile(suffix=file_suffix) as temp:
+            with tempfile.NamedTemporaryFile(mode='w+', suffix=file_suffix) as temp:
                 temp.write(src)
                 temp.flush()
                 subprocess.call([editor, temp.name])
                 temp.seek(0)
                 src = temp.read()
         else:
             print_ansi('#ansi[$EDITOR not set, falling back to stdin](yellow)\n')
-            src = []
+            lines = []
             try:
                 while True:
                     s = input()
                     if s.strip() == ':q':
                         raise EOFError
-                    src.append(s)
+                    lines.append(s)
             except EOFError:  # Ctrl+D
-                src = '\n'.join(src)
+                src = '\n'.join(lines)
             except Exception as io:
                 raise InvalidCommandException(str(io))
         return src
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         pass
 
-    def execute(self, line):
+    def execute(self, line: str) -> Optional[int]:
         raise NotImplementedError
 
 
 commands: Dict[str, Command] = OrderedDict()
 
 
-def register_command(command):
+def register_command(command: Command) -> None:
     commands[command.name] = command

dmoj/commands/diff.py

@@ -1,4 +1,5 @@
 import difflib
+from typing import List
 
 import pygments.formatters
 import pygments.lexers
@@ -10,23 +11,23 @@ class DifferenceCommand(Command):
     name = 'diff'
     help = 'Shows difference between two files.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('id_or_source_1', help='id or path of first source', metavar='<source 1>')
         self.arg_parser.add_argument('id_or_source_2', help='id or path of second source', metavar='<source 2>')
 
-    def get_data(self, id_or_source):
+    def get_data(self, id_or_source: str) -> List[str]:
         try:
             _, _, src, _, _ = self.get_submission_data(int(id_or_source))
         except ValueError:
             src = self.get_source(id_or_source)
 
         return src.splitlines()
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
 
-        file1 = args.id_or_source_1
-        file2 = args.id_or_source_2
+        file1: str = args.id_or_source_1
+        file2: str = args.id_or_source_2
 
         data1 = self.get_data(file1)
         data2 = self.get_data(file2)

dmoj/commands/help.py

@@ -5,10 +5,10 @@ class HelpCommand(Command):
     name = 'help'
     help = 'Prints listing of commands.'
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         print('Run `command -h/--help` for individual command usage.')
         for name, command in commands.items():
             if command == self:
                 continue
-            print('  %s: %s' % (name, command.help))
+            print(f'  {name}: {command.help}')
         print()

dmoj/commands/problems.py

@@ -10,11 +10,11 @@ class ListProblemsCommand(Command):
     name = 'problems'
     help = 'Lists the problems available to be graded on this judge.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('filter', nargs='?', help='regex filter for problem names (optional)')
         self.arg_parser.add_argument('-l', '--limit', type=int, help='limit number of results', metavar='<limit>')
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         _args = self.arg_parser.parse_args(line)
 
         if _args.limit is not None and _args.limit <= 0:
@@ -32,7 +32,7 @@ def execute(self, line):
         if len(all_problems):
             max_len = max(len(p) for p in all_problems)
             for row in zip_longest(*[iter(all_problems)] * 4, fillvalue=''):
-                print(' '.join(('%*s' % (-max_len, row[i])) for i in range(4)))
+                print(' '.join(f'{row[i]:<{max_len}}' for i in range(4)))
             print()
         else:
             raise InvalidCommandException('No problems matching filter found.')

dmoj/commands/quit.py

@@ -7,5 +7,5 @@ class QuitCommand(Command):
     name = 'quit'
     help = 'Exits the DMOJ command-line interface.'
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         sys.exit(0)

dmoj/commands/rejudge.py

@@ -6,10 +6,10 @@ class RejudgeCommand(Command):
     name = 'rejudge'
     help = 'Rejudge a submission.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('submission_id', type=int, help='id of submission to rejudge')
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
         problem_id, lang, src, tl, ml = self.get_submission_data(args.submission_id)
 

dmoj/commands/resubmit.py

@@ -9,7 +9,7 @@ class ResubmitCommand(Command):
     name = 'resubmit'
     help = 'Resubmit a submission with different parameters.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('submission_id', type=int, help='id of submission to resubmit')
         self.arg_parser.add_argument('-p', '--problem', help='id of problem to grade', metavar='<problem id>')
         self.arg_parser.add_argument(
@@ -22,7 +22,7 @@ def _populate_parser(self):
             '-ml', '--memory-limit', type=int, help='memory limit for grading, in kilobytes', metavar='<memory limit>'
         )
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
 
         problem_id, lang, src, tl, ml = self.get_submission_data(args.submission_id)
@@ -33,9 +33,9 @@ def execute(self, line):
         ml = args.memory_limit or ml
 
         if id not in judgeenv.get_supported_problems():
-            raise InvalidCommandException("unknown problem '%s'" % problem_id)
+            raise InvalidCommandException(f"unknown problem '{problem_id}'")
         elif lang not in executors:
-            raise InvalidCommandException("unknown language '%s'" % lang)
+            raise InvalidCommandException(f"unknown language '{lang}'")
         elif tl <= 0:
             raise InvalidCommandException('--time-limit must be >= 0')
         elif ml <= 0:

dmoj/commands/show.py

@@ -1,5 +1,8 @@
+from typing import Tuple
+
 import pygments.formatters
 import pygments.lexers
+from pygments.lexer import Lexer
 
 from dmoj.commands.base_command import Command
 
@@ -8,17 +11,17 @@ class ShowCommand(Command):
     name = 'show'
     help = 'Shows file based on submission ID or filename.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('id_or_source', help='id or path of submission to show', metavar='<source>')
 
-    def get_data(self, id_or_source):
+    def get_data(self, id_or_source: str) -> Tuple[str, Lexer]:
         try:
-            id = int(id_or_source)
+            sub_id = int(id_or_source)
         except ValueError:
             src = self.get_source(id_or_source)
             lexer = pygments.lexers.get_lexer_for_filename(id_or_source)
         else:
-            _, lang, src, _, _ = self.get_submission_data(id)
+            _, lang, src, _, _ = self.get_submission_data(sub_id)
 
             # TODO: after executor->extension mapping is built-in to the judge, redo this
             if lang in ['PY2', 'PYPY2']:
@@ -30,7 +33,7 @@ def get_data(self, id_or_source):
 
         return src, lexer
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
         data, lexer = self.get_data(args.id_or_source)
 

dmoj/commands/submissions.py

@@ -7,12 +7,12 @@ class ListSubmissionsCommand(Command):
     name = 'submissions'
     help = 'List past submissions.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument(
             '-l', '--limit', type=int, help='limit number of results by most recent', metavar='<limit>'
         )
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
 
         if args.limit is not None and args.limit <= 0:
@@ -21,5 +21,5 @@ def execute(self, line):
         submissions = self.judge.graded_submissions if not args.limit else self.judge.graded_submissions[: args.limit]
 
         for i, (problem, lang, src, tl, ml) in enumerate(submissions):
-            print_ansi('#ansi[%s](yellow)/#ansi[%s](green) in %s' % (problem, i + 1, lang))
+            print_ansi(f'#ansi[{problem}](yellow)/#ansi[{i + 1}](green) in {lang}')
         print()

dmoj/commands/submit.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from dmoj import judgeenv
 from dmoj.commands.base_command import Command
 from dmoj.error import InvalidCommandException
@@ -9,7 +11,7 @@ class SubmitCommand(Command):
     name = 'submit'
     help = 'Grades a submission.'
 
-    def _populate_parser(self):
+    def _populate_parser(self) -> None:
         self.arg_parser.add_argument('problem_id', help='id of problem to grade')
         self.arg_parser.add_argument(
             'language_id', nargs='?', default=None, help='id of the language to grade in (e.g., PY2)'
@@ -34,21 +36,21 @@ def _populate_parser(self):
             metavar='<memory limit>',
         )
 
-    def execute(self, line):
+    def execute(self, line: str) -> None:
         args = self.arg_parser.parse_args(line)
 
-        problem_id = args.problem_id
-        language_id = args.language_id
-        time_limit = args.time_limit
-        memory_limit = args.memory_limit
-        source_file = args.source_file
+        problem_id: str = args.problem_id
+        language_id: Optional[str] = args.language_id
+        time_limit: float = args.time_limit
+        memory_limit: int = args.memory_limit
+        source_file: Optional[str] = args.source_file
 
         if language_id not in executors:
             source_file = language_id
             language_id = None  # source file / language id optional
 
         if problem_id not in judgeenv.get_supported_problems():
-            raise InvalidCommandException("unknown problem '%s'" % problem_id)
+            raise InvalidCommandException(f"unknown problem '{problem_id}'")
         elif not language_id:
             if source_file:
                 filename, dot, ext = source_file.partition('.')
@@ -61,12 +63,13 @@ def execute(self, line):
             else:
                 raise InvalidCommandException('no language is selected')
         elif language_id not in executors:
-            raise InvalidCommandException("unknown language '%s'" % language_id)
+            raise InvalidCommandException(f"unknown language '{language_id}'")
         elif time_limit <= 0:
             raise InvalidCommandException('--time-limit must be >= 0')
         elif memory_limit <= 0:
             raise InvalidCommandException('--memory-limit must be >= 0')
 
+        assert language_id is not None
         src = self.get_source(source_file) if source_file else self.open_editor(language_id)
 
         self.judge.submission_id_counter += 1