Add check for None values in set_attributes.

fkdosilovic · fkdosilovic · commit a95d1cc7917c · 2024-11-14T23:51:23.000+01:00
diff --git a/src/judge0/submission.py b/src/judge0/submission.py
@@ -46,6 +46,14 @@
 RESPONSE_FIELDS = ENCODED_RESPONSE_FIELDS | EXTRA_RESPONSE_FIELDS
 FIELDS = REQUEST_FIELDS | RESPONSE_FIELDS
 SKIP_FIELDS = {"language_id", "language", "status_id"}
+DATETIME_FIELDS = {"created_at", "finished_at"}
+FLOATING_POINT_FIELDS = {
+    "cpu_time_limit",
+    "cpu_extra_time",
+    "time",
+    "wall_time",
+    "wall_time_limit",
+}
 
 
 class Submission:
@@ -129,15 +137,9 @@ def set_attributes(self, attributes):
                 value = decode(value) if value else None
             elif attr == "status":
                 value = Status(value["id"])
-            elif attr in ("created_at", "finished_at"):
+            elif attr in DATETIME_FIELDS and value is not None:
                 value = datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%fZ")
-            elif attr in (
-                "wall_time",
-                "wall_time_limit",
-                "cpu_extra_time",
-                "cpu_time_limit",
-                "time",
-            ):
+            elif attr in FLOATING_POINT_FIELDS and value is not None:
                 value = float(value)
 
             setattr(self, attr, value)
@@ -149,6 +151,7 @@ def to_dict(self, client: "Client") -> dict:
             "language_id": client.get_language_id(self.language),
         }
 
+        # TODO: Iterate over ENCODED_REQUEST_FIELDS.
         if self.stdin is not None:
             body["stdin"] = encode(self.stdin)
         if self.expected_output is not None:
