From 6cb8e31186633c5050653421f18d4b8b1c44f8cd Mon Sep 17 00:00:00 2001
From: Tudor Brindus <me@tbrindus.ca>
Date: Sun, 5 Dec 2021 19:33:41 -0500
Subject: [PATCH] executors: add basic autodetection from filename

Ref https://github.com/DMOJ/judge-server/issues/967
---
 dmoj/commands/submit.py             | 15 ++++-----------
 dmoj/executors/__init__.py          | 24 ++++++++++++++++++++++++
 dmoj/utils/helper_files.py          | 22 +++++++++++++++-------
 testsuite/generator_python/init.yml |  4 +---
 4 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/dmoj/commands/submit.py b/dmoj/commands/submit.py
index 105abc58e..4eb7844a7 100644
--- a/dmoj/commands/submit.py
+++ b/dmoj/commands/submit.py
@@ -1,9 +1,8 @@
 from typing import Optional
 
-from dmoj import judgeenv
+from dmoj import executors, judgeenv
 from dmoj.commands.base_command import Command
 from dmoj.error import InvalidCommandException
-from dmoj.executors import executors
 from dmoj.judge import Submission
 
 
@@ -45,7 +44,7 @@ def execute(self, line: str) -> None:
         memory_limit: int = args.memory_limit
         source_file: Optional[str] = args.source_file
 
-        if language_id not in executors:
+        if language_id not in executors.executors:
             source_file = language_id
             language_id = None  # source file / language id optional
 
@@ -53,16 +52,10 @@ def execute(self, line: str) -> None:
             raise InvalidCommandException(f"unknown problem '{problem_id}'")
         elif not language_id:
             if source_file:
-                filename, dot, ext = source_file.partition('.')
-                if not ext:
-                    raise InvalidCommandException('invalid file name')
-                else:
-                    # TODO: this should be a proper lookup elsewhere
-                    ext = ext.upper()
-                    language_id = {'PY': 'PY2', 'CPP': 'CPP11', 'JAVA': 'JAVA8'}.get(ext, ext)
+                language_id = executors.from_filename(source_file).Executor.name
             else:
                 raise InvalidCommandException('no language is selected')
-        elif language_id not in executors:
+        elif language_id not in executors.executors:
             raise InvalidCommandException(f"unknown language '{language_id}'")
         elif time_limit <= 0:
             raise InvalidCommandException('--time-limit must be >= 0')
diff --git a/dmoj/executors/__init__.py b/dmoj/executors/__init__.py
index e5088185f..80dc74d2a 100644
--- a/dmoj/executors/__init__.py
+++ b/dmoj/executors/__init__.py
@@ -15,6 +15,30 @@
 executors: Dict[str, Any] = {}
 
 
+def by_ext(ext: str) -> Any:
+    ext = ext.lower()
+
+    for name, executor in executors.items():
+        if name.lower() == ext:
+            return executor
+
+    for executor in sorted(
+        executors.values(), key=lambda executor: executor.Executor.name
+    ):
+        if executor.Executor.ext == ext:
+            return executor
+
+    raise KeyError('no executor for extension "%s"' % ext)
+
+
+def from_filename(filename: str) -> Any:
+    _, _, ext = filename.partition('.')
+    if not ext:
+        raise KeyError('invalid file name')
+
+    return by_ext(ext)
+
+
 def get_available():
     return get_available_modules(
         _reexecutor, os.path.dirname(__file__), only_executors, exclude_executors | _unsupported_executors
diff --git a/dmoj/utils/helper_files.py b/dmoj/utils/helper_files.py
index 24d8655cb..073f6eb6d 100644
--- a/dmoj/utils/helper_files.py
+++ b/dmoj/utils/helper_files.py
@@ -26,7 +26,7 @@ def compile_with_auxiliary_files(
     compiler_time_limit: Optional[int] = None,
     unbuffered: bool = False,
 ) -> 'BaseExecutor':
-    from dmoj.executors import executors
+    from dmoj import executors
     from dmoj.executors.compiled_executor import CompiledExecutor
 
     sources = {}
@@ -37,21 +37,29 @@ def compile_with_auxiliary_files(
 
     def find_runtime(languages):
         for grader in languages:
-            if grader in executors:
+            if grader in executors.executors:
                 return grader
         return None
 
     use_cpp = any(map(lambda name: os.path.splitext(name)[1] in ['.cpp', '.cc'], filenames))
     use_c = any(map(lambda name: os.path.splitext(name)[1] in ['.c'], filenames))
-    if lang is None:
+    if not lang:
         best_choices = ('CPP20', 'CPP17', 'CPP14', 'CPP11', 'CPP03') if use_cpp else ('C11', 'C')
         lang = find_runtime(best_choices)
 
-    executor = executors.get(lang)
-    if not executor:
-        raise IOError('could not find an appropriate C++ executor')
+    # TODO: remove above code once `from_filename` is smart enough to
+    # prioritize newer versions of runtimes
+    if not lang:
+        for filename in filenames:
+            try:
+                lang = executors.from_filename(filename).Executor.name
+            except KeyError:
+                continue
 
-    executor = executor.Executor
+    if not lang:
+        raise IOError('could not find an appropriate executor')
+
+    executor = executors.executors[lang].Executor
 
     kwargs = {'fs': executor.fs + [RecursiveDir(tempfile.gettempdir())]}
 
diff --git a/testsuite/generator_python/init.yml b/testsuite/generator_python/init.yml
index 5b2c97a52..ff5b3e15b 100644
--- a/testsuite/generator_python/init.yml
+++ b/testsuite/generator_python/init.yml
@@ -1,6 +1,4 @@
-generator:
-    language: PY2
-    source: gen.py
+generator: gen.py2
 test_cases:
 - {generator_args: [1, 2], points: 1}
 - {generator_args: [10, 20], points: 1}