support vmlinux enum in printk handler

Author: varun-r-mallya

pythonbpf/helper/printk_formatter.py

@@ -3,6 +3,7 @@
 
 from llvmlite import ir
 from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
+from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
 
 logger = logging.getLogger(__name__)
 
@@ -108,6 +109,16 @@ def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
     if local_sym_tab and name_node.id in local_sym_tab:
         _, var_type, tmp = local_sym_tab[name_node.id]
         _populate_fval(var_type, name_node, fmt_parts, exprs)
+    else:
+        # Try to resolve through vmlinux registry if not in local symbol table
+        result = VmlinuxHandlerRegistry.handle_name(name_node.id)
+        if result:
+            val, var_type = result
+            _populate_fval(var_type, name_node, fmt_parts, exprs)
+        else:
+            raise ValueError(
+                f"Variable '{name_node.id}' not found in symbol table or vmlinux"
+            )
 
 
 def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):

tests/passing_tests/vmlinux/simple_struct_test.py

@@ -1,4 +1,7 @@
-from pythonbpf import bpf, section, bpfglobal, compile_to_ir, compile
+import logging
+
+from pythonbpf import bpf, section, bpfglobal, compile_to_ir
+from pythonbpf import compile  # noqa: F401
 from vmlinux import TASK_COMM_LEN  # noqa: F401
 from vmlinux import struct_trace_event_raw_sys_enter  # noqa: F401
 
@@ -17,7 +20,7 @@
 @section("tracepoint/syscalls/sys_enter_execve")
 def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64:
     a = 2 + TASK_COMM_LEN + TASK_COMM_LEN
-    print(f"Hello, World{a}")
+    print(f"Hello, World{TASK_COMM_LEN} and {a}")
     return c_int64(TASK_COMM_LEN)
 
 
@@ -27,5 +30,5 @@ def LICENSE() -> str:
     return "GPL"
 
 
-compile_to_ir("simple_struct_test.py", "simple_struct_test.ll")
+compile_to_ir("simple_struct_test.py", "simple_struct_test.ll", loglevel=logging.DEBUG)
 compile()