Refactor parse_return_statement method to handle nested returns (#1549)

Fixes [Component Output] Custom component does not display output connection points in Langflow UI #1357
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-21 12:21:44 -03:00 • committed by GitHub
commit 3fbcd02bd9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -6,6 +6,8 @@ from typing import Any, Dict, List, Type, Union
from cachetools import TTLCache, cachedmethod, keys from cachetools import TTLCache, cachedmethod, keys
from fastapi import HTTPException from fastapi import HTTPException
from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails
@ -234,10 +236,28 @@ class CodeParser:
def parse_return_statement(self, node: ast.FunctionDef) -> bool: def parse_return_statement(self, node: ast.FunctionDef) -> bool:
""" """
Parses the return statement of a function or method node. Parses the return statement of a function or method node, including nested returns.
""" """
return any(isinstance(n, ast.Return) for n in node.body) 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): def parse_assign(self, stmt):
""" """