Load implicit and preview clients only when needed. (judge0#1)

hermanzdosilovic · fkdosilovic · web-flow · commit 967ed92c5484 · 2024-11-12T21:16:48.000+01:00
* Load implicit and preview clients only when needed.

* Add comments. Make api_key for Sulu clients positional and keyword argument with default value.

* Run formatter and minor refactor of api module.

---------

Co-authored-by: Filip Karlo Došilović &lt;filipk.dosilovic@gmail.com&gt;
diff --git a/src/judge0/__init__.py b/src/judge0/__init__.py
@@ -43,11 +43,19 @@
     "wait",
 ]
 
+JUDGE0_IMPLICIT_CE_CLIENT = None
+JUDGE0_IMPLICIT_EXTRA_CE_CLIENT = None
+
+
+def _get_implicit_client(flavor: Flavor) -> Optional[Client]:
+    global JUDGE0_IMPLICIT_CE_CLIENT, JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
+
+    # Implicit clients are already set.
+    if flavor == Flavor.CE and JUDGE0_IMPLICIT_CE_CLIENT is not None:
+        return JUDGE0_IMPLICIT_CE_CLIENT
+    if flavor == Flavor.EXTRA_CE and JUDGE0_IMPLICIT_EXTRA_CE_CLIENT is not None:
+        return JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
 
-def _create_default_client(
-    flavor: Flavor,
-    preview_client_class: Union[SuluJudge0CE, SuluJudge0ExtraCE],
-) -> Optional[Client]:
     from .clients import CE, EXTRA_CE
 
     try:
@@ -62,31 +70,31 @@ def _create_default_client(
     else:
         client_classes = EXTRA_CE
 
+    # Try to find one of the predefined keys JUDGE0_{SULU,RAPID,ATD}_API_KEY
+    # in environment variables.
     client = None
     for client_class in client_classes:
         api_key = os.getenv(client_class.API_KEY_ENV)
         if api_key is not None:
-            # It is possible for a client to be subscribed to one flavor with
-            # API key.
-            try:
-                client = client_class(api_key)
-                break
-            except Exception as e:
-                warnings.warn(f"Failed to initialize client: {e}")
+            client = client_class(api_key)
+            break
+
+    # If we didn't find any of the possible predefined keys, initialize
+    # the preview Sulu client based on the flavor.
+    if client is None:
+        if flavor == Flavor.CE:
+            client = SuluJudge0CE()
+        else:
+            client = SuluJudge0ExtraCE()
+
+    if flavor == Flavor.CE:
+        JUDGE0_IMPLICIT_CE_CLIENT = client
     else:
-        try:
-            client = preview_client_class("")
-        except Exception as e:
-            warnings.warn(f"Failed to initialize preview client: {e}")
+        JUDGE0_IMPLICIT_EXTRA_CE_CLIENT = client
 
     return client
 
 
-JUDGE0_IMPLICIT_CE_CLIENT = _create_default_client(Flavor.CE, SuluJudge0CE)
-JUDGE0_IMPLICIT_EXTRA_CE_CLIENT = _create_default_client(
-    Flavor.EXTRA_CE, SuluJudge0ExtraCE
-)
-
 CE = Flavor.CE
 EXTRA_CE = Flavor.EXTRA_CE
 
diff --git a/src/judge0/api.py b/src/judge0/api.py
@@ -10,44 +10,30 @@ def resolve_client(
     client: Optional[Union[Client, Flavor]] = None,
     submissions: Optional[Union[Submission, list[Submission]]] = None,
 ) -> Union[Client, None]:
-    from . import JUDGE0_IMPLICIT_CE_CLIENT, JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
-
     # User explicitly passed a client.
     if isinstance(client, Client):
         return client
 
+    from . import _get_implicit_client
+
     # User explicitly choose the flavor of the client.
     if isinstance(client, Flavor):
-        if client == Flavor.CE:
-            return JUDGE0_IMPLICIT_CE_CLIENT
-        else:
-            return JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
+        return _get_implicit_client(flavor=client)
 
     # client is None and we have to determine a flavor of the client from the
     # submissions and the languages.
     if isinstance(submissions, Submission):
         submissions = [submissions]
 
     # Check which client supports all languages from the provided submissions.
-    languages = [submission.language_id for submission in submissions]
-
-    if JUDGE0_IMPLICIT_CE_CLIENT is not None:
-        if all(
-            (
-                JUDGE0_IMPLICIT_CE_CLIENT.is_language_supported(lang)
-                for lang in languages
-            )
-        ):
-            return JUDGE0_IMPLICIT_CE_CLIENT
+    languages = (submission.language_id for submission in submissions)
 
-    if JUDGE0_IMPLICIT_EXTRA_CE_CLIENT is not None:
-        if all(
-            (
-                JUDGE0_IMPLICIT_EXTRA_CE_CLIENT.is_language_supported(lang)
-                for lang in languages
-            )
+    for flavor in Flavor:
+        client = _get_implicit_client(flavor)
+        if client is not None and all(
+            (client.is_language_supported(lang) for lang in languages)
         ):
-            return JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
+            return client
 
     raise RuntimeError(
         "Failed to resolve the client from submissions argument. "
@@ -65,13 +51,13 @@ def wait(
     if retry_mechanism is None:
         retry_mechanism = RegularPeriodRetry()
 
-    if not isinstance(submissions, (list, tuple)):
+    if isinstance(submissions, (list, tuple)):
         submissions_to_check = {
-            submission.token: submission for submission in [submissions]
+            submission.token: submission for submission in submissions
         }
     else:
         submissions_to_check = {
-            submission.token: submission for submission in submissions
+            submission.token: submission for submission in [submissions]
         }
 
     while len(submissions_to_check) > 0 and not retry_mechanism.is_done():
diff --git a/src/judge0/clients.py b/src/judge0/clients.py
@@ -402,25 +402,28 @@ def __init__(self, api_key):
 class Sulu(Client):
     API_KEY_ENV = "JUDGE0_SULU_API_KEY"
 
-    def __init__(self, endpoint, api_key):
+    def __init__(self, endpoint, api_key=None):
         self.api_key = api_key
-        super().__init__(endpoint, {"Authorization": f"Bearer {api_key}"})
+        super().__init__(
+            endpoint,
+            {"Authorization": f"Bearer {api_key}"} if api_key else None,
+        )
 
 
 class SuluJudge0CE(Sulu):
     DEFAULT_ENDPOINT: str = "https://judge0-ce.p.sulu.sh"
     HOME_URL: str = "https://sparkhub.sulu.sh/apis/judge0/judge0-ce/readme"
 
-    def __init__(self, api_key):
-        super().__init__(self.DEFAULT_ENDPOINT, api_key=api_key)
+    def __init__(self, api_key=None):
+        super().__init__(self.DEFAULT_ENDPOINT, api_key)
 
 
 class SuluJudge0ExtraCE(Sulu):
     DEFAULT_ENDPOINT: str = "https://judge0-extra-ce.p.sulu.sh"
     HOME_URL: str = "https://sparkhub.sulu.sh/apis/judge0/judge0-extra-ce/readme"
 
-    def __init__(self, api_key):
-        super().__init__(self.DEFAULT_ENDPOINT, api_key=api_key)
+    def __init__(self, api_key=None):
+        super().__init__(self.DEFAULT_ENDPOINT, api_key)
 
 
 CE = [RapidJudge0CE, SuluJudge0CE, ATDJudge0CE]
