From e55c42113eb3f04dbebfba39193ce8eae91c71cc 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          | 22 ++++++++++++++++++++++
 dmoj/problem.py                     |  7 +++++--
 testsuite/generator_python/init.yml |  4 +---
 4 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/dmoj/commands/submit.py b/dmoj/commands/submit.py
index 105abc58e..5762be28e 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).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..30ecfb98b 100644
--- a/dmoj/executors/__init__.py
+++ b/dmoj/executors/__init__.py
@@ -15,6 +15,28 @@
 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 executors.values():
+        if 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/problem.py b/dmoj/problem.py
index 794099656..148c27922 100644
--- a/dmoj/problem.py
+++ b/dmoj/problem.py
@@ -10,7 +10,7 @@
 from yaml.parser import ParserError
 from yaml.scanner import ScannerError
 
-from dmoj import checkers
+from dmoj import checkers, executors
 from dmoj.config import ConfigNode, InvalidInitException
 from dmoj.judgeenv import env, get_problem_root
 from dmoj.utils.helper_files import compile_with_auxiliary_files, parse_helper_file_error
@@ -272,7 +272,7 @@ def _run_generator(self, gen, args=None):
         time_limit = env.generator_time_limit
         memory_limit = env.generator_memory_limit
         compiler_time_limit = env.generator_compiler_time_limit
-        lang = None  # Default to C/C++
+        lang = None
 
         base = get_problem_root(self.problem.id)
         if isinstance(gen, str):
@@ -300,6 +300,9 @@ def _run_generator(self, gen, args=None):
         if not isinstance(filenames, list):
             filenames = [filenames]
 
+        if not lang:
+            lang = executors.from_filename(filenames[0])
+
         filenames = [os.path.abspath(os.path.join(base, name)) for name in filenames]
         executor = compile_with_auxiliary_files(filenames, flags, lang, compiler_time_limit)
 
diff --git a/testsuite/generator_python/init.yml b/testsuite/generator_python/init.yml
index 5b2c97a52..b089fe7a6 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.py
 test_cases:
 - {generator_args: [1, 2], points: 1}
 - {generator_args: [10, 20], points: 1}