Add _is_xdp_name

Author: r41k0u

pythonbpf/functions/functions_pass.py

@@ -9,7 +9,7 @@
 from pythonbpf.binary_ops import handle_binary_op
 from pythonbpf.expr_pass import eval_expr, handle_expr
 
-from .return_utils import _handle_none_return, _handle_xdp_return
+from .return_utils import _handle_none_return, _handle_xdp_return, _is_xdp_name
 
 
 logger = logging.getLogger(__name__)
@@ -357,7 +357,7 @@ def handle_return(builder, stmt, local_sym_tab, ret_type):
     logger.info(f"Handling return statement: {ast.dump(stmt)}")
     if stmt.value is None:
         return _handle_none_return(builder)
-    elif isinstance(stmt.value, ast.Name):
+    elif isinstance(stmt.value, ast.Name) and _is_xdp_name(stmt.value.id):
         return _handle_xdp_return(stmt, builder, ret_type)
     elif (
         isinstance(stmt.value, ast.Call)

pythonbpf/functions/return_utils.py

@@ -111,6 +111,11 @@ def _handle_wrapped_return(stmt: ast.Return, builder, ret_type, local_sym_tab) -
         raise ValueError(f"Unsupported return argument type: {type(arg).__name__}")
 
 
+def _is_xdp_name(name: str) -> bool:
+    """Check if a name is an XDP action"""
+    return name in XDP_ACTIONS
+
+
 def _handle_xdp_return(stmt: ast.Return, builder, ret_type) -> bool:
     """Handle XDP returns"""
     if not isinstance(stmt.value, ast.Name):