Patched patchwork/common/utils/step_typing.py

Author: patched.codes[bot]

patchwork/common/utils/step_typing.py

@@ -106,8 +106,13 @@ def validate_step_type_config_with_inputs(
 
 
 def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[Set[str], Dict[str, str]]:
+    ALLOWED_MODULES = {"module1.typed", "module2.typed"}  # Define allowed modules
     module_path, _, _ = step.__module__.rpartition(".")
     step_name = step.__name__
+    
+    if f"{module_path}.typed" not in ALLOWED_MODULES:
+        raise ValueError("Attempted to import a module not in the allowed list.")
+    
     type_module = importlib.import_module(f"{module_path}.typed")
     step_input_model = getattr(type_module, f"{step_name}Inputs", __NOT_GIVEN)
     step_output_model = getattr(type_module, f"{step_name}Outputs", __NOT_GIVEN)