Update rp_validator.py

justinmerrell · justinmerrell · commit b8d48e7bf8f7 · 2023-04-23T16:00:12.000-04:00
diff --git a/runpod/serverless/utils/rp_validator.py b/runpod/serverless/utils/rp_validator.py
@@ -3,8 +3,19 @@
 Provides a function to validate the input to the model.
 '''
 # pylint: disable=too-many-branches
+
+import json
 from typing import Any, Dict, List, Union
 
+# Error messages
+UNEXPECTED_INPUT_ERROR = "Unexpected input. {} is not a valid input option."
+MISSING_REQUIRED_ERROR = "{} is a required input."
+MISSING_DEFAULT_ERROR = "Schema error, missing default value for {}."
+MISSING_TYPE_ERROR = "Schema error, missing type for {}."
+INVALID_TYPE_ERROR = "{} should be {} type, not {}."
+CONSTRAINTS_ERROR = "{} does not meet the constraints."
+SCHEMA_ERROR = "Schema error, {} is not a dictionary."
+
 
 def _add_error(error_list: List[str], message: str) -> None:
     error_list.append(message)
@@ -25,29 +36,41 @@ def validate(raw_input: Dict[str, Any], schema: Dict[str, Any]
     {"validated_input": {"input1": "value1", "input2": "value2"}
     '''
     error_list = []
+    validated_input = raw_input.copy()
 
     # Check for unexpected inputs.
     for key in raw_input:
         if key not in schema:
-            _add_error(error_list, f"Unexpected input. {key} is not a valid input option.")
+            _add_error(error_list, UNEXPECTED_INPUT_ERROR.format(key))
+
+    # Check that items are dictionaries.
+    for key, rules in schema.items():
+        if not isinstance(rules, dict):
+            try:
+                schema[key] = json.loads(rules)
+            except json.decoder.JSONDecodeError:
+                _add_error(error_list, SCHEMA_ERROR.format(key))
 
     # Checks for missing required inputs or sets the default values.
     for key, rules in schema.items():
+        if 'type' not in rules:
+            _add_error(error_list, MISSING_TYPE_ERROR.format(key))
+
         if 'required' not in rules:
-            _add_error(error_list, f"Schema error, missing 'required' for {key}.")
+            _add_error(error_list, MISSING_REQUIRED_ERROR.format(key))
         elif rules['required'] and key not in raw_input:
-            _add_error(error_list, f"{key} is a required input.")
+            _add_error(error_list, MISSING_REQUIRED_ERROR.format(key))
         elif rules['required'] and key not in raw_input and "default" not in rules:
-            _add_error(error_list, f"Schema error, missing default value for {key}.")
+            _add_error(error_list, MISSING_DEFAULT_ERROR.format(key))
         elif not rules['required'] and key not in raw_input and "default" not in rules:
-            _add_error(error_list, f"Schema error, missing default value for {key}.")
+            _add_error(error_list, MISSING_DEFAULT_ERROR.format(key))
         elif not rules['required'] and key not in raw_input:
-            raw_input[key] = raw_input.get(key, rules['default'])
+            validated_input[key] = raw_input.get(key, rules['default'])
 
     for key, rules in schema.items():
         # Enforce floats to be floats.
         if rules['type'] is float and type(raw_input[key]) in [int, float]:
-            raw_input[key] = float(raw_input[key])
+            validated_input[key] = float(raw_input[key])
 
         # Check for the correct type.
         if not isinstance(raw_input[key], rules['type']) and raw_input[key] is not None:
@@ -57,9 +80,9 @@ def validate(raw_input: Dict[str, Any], schema: Dict[str, Any]
         # Check lambda constraints.
         if "constraints" in rules:
             if not rules['constraints'](raw_input[key]):
-                _add_error(error_list, f"{key} does not meet the constraints.")
+                _add_error(error_list, CONSTRAINTS_ERROR.format(key))
 
-    validation_return = {"validated_input": raw_input}
+    validation_return = {"validated_input": validated_input}
     if error_list:
         validation_return = {"errors": error_list}
 
