From da82287b7564e51292808eb2fc17d9e3ed9a800f Mon Sep 17 00:00:00 2001
From: Bob Yang <bobyf@meta.com>
Date: Tue, 25 Mar 2025 13:08:13 -0700
Subject: [PATCH] Validation for pipeline and appdef components (#1024)

Summary: Pull Request resolved: https://github.com/pytorch/torchx/pull/1024

Reviewed By: lgarg26

Differential Revision: D71297049
---
 torchx/specs/file_linter.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/torchx/specs/file_linter.py b/torchx/specs/file_linter.py
index 90b1df064..b11abd78f 100644
--- a/torchx/specs/file_linter.py
+++ b/torchx/specs/file_linter.py
@@ -180,6 +180,11 @@ def _validate_arg_def(
 
 
 class TorchxReturnValidator(TorchxFunctionValidator):
+
+    def __init__(self, supported_return_type: str) -> None:
+        super().__init__()
+        self._supported_return_type = supported_return_type
+
     def _get_return_annotation(
         self, app_specs_func_def: ast.FunctionDef
     ) -> Optional[str]:
@@ -203,7 +208,7 @@ def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
             * AppDef
             * specs.AppDef
         """
-        supported_return_annotation = "AppDef"
+        supported_return_annotation = self._supported_return_type
         return_annotation = self._get_return_annotation(app_specs_func_def)
         linter_errors = []
         if not return_annotation:
@@ -252,7 +257,7 @@ def __init__(
         if validators is None:
             self.validators: List[TorchxFunctionValidator] = [
                 TorchxFunctionArgsValidator(),
-                TorchxReturnValidator(),
+                TorchxReturnValidator("AppDef"),
             ]
         else:
             self.validators = validators