Fix bugs with globs

* Order globs, so an earlier-listed problem root takes priority
* Check that the return value of `get_problem_root` satisfies at least
  one glob
* Add tests for globs

dmoj/judgeenv.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import ssl
+from fnmatch import fnmatch
 from operator import itemgetter
 from typing import Dict, List, Set
 
@@ -250,8 +251,12 @@ def get_problem_root(problem_id):
     if cached_root is None or not os.path.isfile(os.path.join(cached_root, 'init.yml')):
         for root_dir in get_problem_roots():
             problem_root_dir = os.path.join(root_dir, problem_id)
-            problem_init = os.path.join(problem_root_dir, 'init.yml')
-            if os.path.isfile(problem_init):
+            problem_config = os.path.join(problem_root_dir, 'init.yml')
+            if os.path.isfile(problem_config):
+                if problem_globs and not any(
+                    fnmatch(problem_config, os.path.join(problem_glob, 'init.yml')) for problem_glob in problem_globs
+                ):
+                    continue
                 _problem_root_cache[problem_id] = problem_root_dir
                 break
 
@@ -268,16 +273,15 @@ def get_problem_roots(warnings=False):
         return _problem_dirs_cache
 
     if problem_globs:
-        dirs = set()
+        dirs = []
+        dirs_set = set()
         for dir_glob in problem_globs:
             config_glob = os.path.join(dir_glob, 'init.yml')
-            dirs = dirs.union(
-                map(
-                    lambda x: os.path.split(os.path.dirname(x))[0],
-                    glob.iglob(config_glob, recursive=True),
-                )
-            )
-        dirs = list(dirs)
+            root_dirs = {os.path.dirname(os.path.dirname(x)) for x in glob.iglob(config_glob, recursive=True)}
+            for root_dir in root_dirs:
+                if root_dir not in dirs_set:
+                    dirs.append(root_dir)
+                    dirs_set.add(root_dir)
     else:
 
         def get_path(x, y):
@@ -347,13 +351,22 @@ def get_supported_problems_and_mtimes():
         A list of all problems in tuple format: (problem id, mtime)
     """
     problems = []
-    for dir in get_problem_roots():
-        if not os.path.isdir(dir):  # we do this check in case a problem root was deleted but persists in cache
-            continue
-        for problem in os.listdir(dir):
-            problem = utf8text(problem)
-            if os.access(os.path.join(dir, problem, 'init.yml'), os.R_OK):
-                problems.append((problem, os.path.getmtime(os.path.join(dir, problem))))
+    if problem_globs:
+        for dir_glob in problem_globs:
+            for problem_config in glob.iglob(os.path.join(dir_glob, 'init.yml')):
+                if os.access(problem_config, os.R_OK):
+                    problem_dir = os.path.dirname(problem_config)
+                    problem = utf8text(os.path.basename(problem_dir))
+                    problems.append((problem, os.path.getmtime(problem_dir)))
+
+    else:
+        for dir in get_problem_roots():
+            if not os.path.isdir(dir):  # we do this check in case a problem root was deleted but persists in cache
+                continue
+            for problem in os.listdir(dir):
+                problem = utf8text(problem)
+                if os.access(os.path.join(dir, problem, 'init.yml'), os.R_OK):
+                    problems.append((problem, os.path.getmtime(os.path.join(dir, problem))))
     return problems
 
 

dmoj/tests/test_globs.py

@@ -0,0 +1,93 @@
+# type: ignore
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+from dmoj import judgeenv
+from dmoj.utils.glob_ext import find_glob_root
+
+
+class TestGlobExt(unittest.TestCase):
+    def test_find_glob_root(self):
+        self.assertEqual(find_glob_root('/a'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/b/c'), Path('/a/b/c'))
+        self.assertEqual(find_glob_root('/a/*'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/*/b'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/*/b/*'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/**/b/*'), Path('/a'))
+        self.assertEqual(find_glob_root('/a/b/**/c/*'), Path('/a/b'))
+        self.assertEqual(find_glob_root('/a/b/*'), Path('/a/b'))
+        self.assertEqual(find_glob_root('/a/b/c[1-5]'), Path('/a/b'))
+        self.assertEqual(find_glob_root('/a/b/c?'), Path('/a/b'))
+        self.assertEqual(find_glob_root('/a/b/c?/*'), Path('/a/b'))
+
+
+class TestConfigGlobs(unittest.TestCase):
+    def setUp(self):
+        self.root = tempfile.TemporaryDirectory()
+        self.root_path = Path(self.root.name)
+
+        dirs = [
+            self.root_path / 'ch1' / 'p1',
+            self.root_path / 'ch1' / 'p2',
+            self.root_path / 'ch2' / 'p3',
+            self.root_path / 'ch2' / 'p4',
+            self.root_path / 'ch3' / 'ch4' / 'p1',
+            self.root_path / 'ch3' / 'ch5' / 'ch6' / 'p5',
+            self.root_path / 'ch3' / 'ch5' / 'p6',
+        ]
+        for dir in dirs:
+            dir.mkdir(parents=True)
+            (dir / 'init.yml').touch()  # make init.yml
+
+        self.problem_roots = list(
+            map(
+                str,
+                (
+                    self.root_path / 'ch1',
+                    self.root_path / 'ch2',
+                    self.root_path / 'ch3' / 'ch4',
+                    self.root_path / 'ch3' / 'ch5',
+                    self.root_path / 'ch3' / 'ch5' / 'ch6',
+                ),
+            )
+        )
+
+        problem_globs = list(
+            map(
+                str,
+                (
+                    self.root_path / 'ch1' / 'p[13]',
+                    self.root_path / 'ch2' / 'p[24]',
+                    self.root_path / 'ch3' / '**',
+                    self.root_path / 'ch7' / '**',
+                ),
+            )
+        )
+
+        self.mock_problem_roots = mock.patch.object(judgeenv, 'problem_globs', problem_globs)
+
+    def test_problem_roots(self):
+        with self.mock_problem_roots:
+            problem_roots = judgeenv.get_problem_roots()
+            self.assertEqual(list(sorted(self.problem_roots)), list(sorted(problem_roots)))
+
+            cases = [
+                (self.root_path / 'ch1' / 'p1', 'p1'),
+                (self.root_path / 'ch2' / 'p4', 'p4'),
+                (self.root_path / 'ch3' / 'ch5' / 'ch6' / 'p5', 'p5'),
+                (self.root_path / 'ch3' / 'ch5' / 'p6', 'p6'),
+            ]
+
+            for path, problem in cases:
+                self.assertEqual(str(path), judgeenv.get_problem_root(problem))
+
+            ex_cases = ['p2', 'p3', 'doesnotexist']
+
+            for problem in ex_cases:
+                self.assertRaises(KeyError, judgeenv.get_problem_root, problem)
+
+    def tearDown(self):
+        self.root.cleanup()