From 3737e06dc43a2745ae0511016dfb91f55271e5ae Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 21 Mar 2024 12:20:41 -0300 Subject: [PATCH] Refactor parse_return_statement method to handle nested returns Fixes [Component Output] Custom component does not display output connection points in Langflow UI #1357 --- .../custom/code_parser/code_parser.py | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/interface/custom/code_parser/code_parser.py b/src/backend/langflow/interface/custom/code_parser/code_parser.py index 7a102e33f0..1ab1021b75 100644 --- a/src/backend/langflow/interface/custom/code_parser/code_parser.py +++ b/src/backend/langflow/interface/custom/code_parser/code_parser.py @@ -6,6 +6,8 @@ from cachetools import TTLCache, cachedmethod, keys from fastapi import HTTPException + + from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails @@ -234,10 +236,28 @@ def parse_function_body(self, node: ast.FunctionDef) -> List[str]: def parse_return_statement(self, node: ast.FunctionDef) -> bool: """ - Parses the return statement of a function or method node. - """ - - return any(isinstance(n, ast.Return) for n in node.body) + Parses the return statement of a function or method node, including nested returns. + """ + + def has_return(node): + if isinstance(node, ast.Return): + return True + elif isinstance(node, ast.If): + return any(has_return(child) for child in node.body) or any(has_return(child) for child in node.orelse) + elif isinstance(node, ast.Try): + return ( + any(has_return(child) for child in node.body) + or any(has_return(child) for child in node.handlers) + or any(has_return(child) for child in node.finalbody) + ) + elif isinstance(node, (ast.For, ast.While)): + return any(has_return(child) for child in node.body) or any(has_return(child) for child in node.orelse) + elif isinstance(node, ast.With): + return any(has_return(child) for child in node.body) + else: + return False + + return any(has_return(child) for child in node.body) def parse_assign(self, stmt): """