Implement probe string extraction

pythonbpf/codegen.py

@@ -1,7 +1,7 @@
 import ast
 from llvmlite import ir
 from .license_pass import license_processing
-from .functions_pass import functions_processing
+from .functions_pass import func_proc, functions_processing
 from .constants_pass import constants_processing
 from .globals_pass import globals_processing
 
@@ -26,6 +26,7 @@ def processor(source_code, filename, module):
     for func_node in bpf_chunks:
         print(f"Found BPF function: {func_node.name}")
 
+    func_proc(tree, module, bpf_chunks)
     # For now, we will parse the BPF specific parts of AST
     # Big rewrite
 

pythonbpf/functions_pass.py

@@ -43,6 +43,27 @@ def emit_function(module: ir.Module, name: str):
     return func
 
 
+def get_probe_string(func_node):
+    """Extract the probe string from the decorator of the function node."""
+    # TODO: right now we have the whole string in the section decorator
+    # But later we can implement typed tuples for tracepoints and kprobes
+    # For helper functions, we return "helper"
+
+    for decorator in func_node.decorator_list:
+        if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name):
+            if decorator.func.id == "section" and len(decorator.args) == 1:
+                arg = decorator.args[0]
+                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                    return arg.value
+    return "helper"
+
+
+def func_proc(tree, module, chunks):
+    for func_node in chunks:
+        func_type = get_probe_string(func_node)
+        print(f"Found probe_string of {func_node.name}: {func_type}")
+
+
 def functions_processing(tree, module):
     bpf_functions = []
     helper_functions = []