diff --git a/dmoj/commands/submit.py b/dmoj/commands/submit.py
index 105abc58e..a7a2dcfcf 100644
--- a/dmoj/commands/submit.py
+++ b/dmoj/commands/submit.py
@@ -53,13 +53,7 @@ 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:
diff --git a/dmoj/executors/__init__.py b/dmoj/executors/__init__.py
index e5088185f..6f038aa92 100644
--- a/dmoj/executors/__init__.py
+++ b/dmoj/executors/__init__.py
@@ -15,6 +15,28 @@
 executors: Dict[str, Any] = {}
 
 
+def by_ext(ext : str):
+    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
+
+    return KeyError('no executor for extension "%s"' % ext)
+
+
+def from_filename(filename : str):
+    _, _, 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}