Remove backslash for multiline conditions in maps

examples/execve2.py

@@ -1,7 +1,7 @@
 from pythonbpf.decorators import bpf, map, section, bpfglobal
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
 from pythonbpf.helpers import ktime
-from pythonbpf.maps.maps import HashMap
+from pythonbpf.maps import HashMap
 
 
 @bpf

examples/execve3.py

@@ -1,6 +1,6 @@
 from pythonbpf import bpf, map, section, bpfglobal, compile
 from pythonbpf.helpers import ktime, deref
-from pythonbpf.maps.maps import HashMap
+from pythonbpf.maps import HashMap
 
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
 

pythonbpf/codegen.py

@@ -2,7 +2,7 @@ import ast
 from llvmlite import ir
 from .license_pass import license_processing
 from .functions_pass import func_proc
-from pythonbpf.maps.maps_pass import maps_proc
+from pythonbpf.maps import maps_proc
 from .structs.structs_pass import structs_proc
 from .globals_pass import globals_processing
 import os

pythonbpf/maps/__init__.py

@@ -0,0 +1,2 @@
+from .maps import HashMap, PerfEventArray
+from .maps_pass import maps_proc

pythonbpf/maps/maps_pass.py

@@ -1,18 +1,20 @@
 import ast
 from llvmlite import ir
-from pythonbpf.type_deducer import ctypes_to_ir
 from pythonbpf import dwarf_constants as dc
 from enum import Enum
+from .maps_utils import MapProcessorRegistry
+import logging
 
-map_sym_tab = {}
+logger = logging.getLogger(__name__)
 
 
 def maps_proc(tree, module, chunks):
+    """ Process all functions decorated with @map to find BPF maps """
+    map_sym_tab = {}
     for func_node in chunks:
         if is_map(func_node):
             print(f"Found BPF map: {func_node.name}")
-            process_bpf_map(func_node, module)
+            map_sym_tab[func_node.name] = process_bpf_map(func_node, module)
-            continue
     return map_sym_tab
 
 
@@ -29,9 +31,7 @@ class BPFMapType(Enum):
 
 
 def create_bpf_map(module, map_name, map_params):
-    """Create a BPF map in the module with the given parameters and debug info"""
+    """Create a BPF map in the module with given parameters and debug info"""
-
-    map_type = map_params.get("type", BPFMapType.HASH).value
 
     # Create the anonymous struct type for BPF map
     map_struct_type = ir.LiteralStructType(
@@ -42,15 +42,14 @@ def create_bpf_map(module, map_name, map_params):
     map_global.linkage = 'dso_local'
     map_global.global_constant = False
     map_global.initializer = ir.Constant(
-        map_struct_type, None)     # type: ignore
+        map_struct_type, None)
     map_global.section = ".maps"
-    map_global.align = 8        # type: ignore
+    map_global.align = 8
 
     # Generate debug info for BTF
     create_map_debug_info(module, map_global, map_name, map_params)
 
-    print(f"Created BPF map: {map_name}")
+    logger.info(f"Created BPF map: {map_name} with params {map_params}")
-    map_sym_tab[map_name] = map_global
     return map_global
 
 
@@ -183,8 +182,10 @@ def create_map_debug_info(module, map_global, map_name, map_params):
     return global_var_expr
 
 
+@MapProcessorRegistry.register("HashMap")
 def process_hash_map(map_name, rval, module):
-    print(f"Creating HashMap map: {map_name}")
+    """Process a BPF_HASH map declaration"""
+    logger.info(f"Processing HashMap: {map_name}")
     map_params = {"type": BPFMapType.HASH}
 
     # Assuming order: key_type, value_type, max_entries
@@ -202,17 +203,20 @@ def process_hash_map(map_name, rval, module):
             map_params["key"] = keyword.value.id
         elif keyword.arg == "value" and isinstance(keyword.value, ast.Name):
             map_params["value"] = keyword.value.id
-        elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant):
+        elif (keyword.arg == "max_entries" and
+              isinstance(keyword.value, ast.Constant)):
             const_val = keyword.value.value
             if isinstance(const_val, (int, str)):
                 map_params["max_entries"] = const_val
 
-    print(f"Map parameters: {map_params}")
+    logger.info(f"Map parameters: {map_params}")
     return create_bpf_map(module, map_name, map_params)
 
 
+@MapProcessorRegistry.register("PerfEventArray")
 def process_perf_event_map(map_name, rval, module):
-    print(f"Creating PerfEventArray map: {map_name}")
+    """Process a BPF_PERF_EVENT_ARRAY map declaration"""
+    logger.info(f"Processing PerfEventArray: {map_name}")
     map_params = {"type": BPFMapType.PERF_EVENT_ARRAY}
 
     if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name):
@@ -223,21 +227,18 @@ def process_perf_event_map(map_name, rval, module):
     for keyword in rval.keywords:
         if keyword.arg == "key_size" and isinstance(keyword.value, ast.Name):
             map_params["key_size"] = keyword.value.id
-        elif keyword.arg == "value_size" and isinstance(keyword.value, ast.Name):
+        elif (keyword.arg == "value_size" and
+              isinstance(keyword.value, ast.Name)):
             map_params["value_size"] = keyword.value.id
 
-    print(f"Map parameters: {map_params}")
+    logger.info(f"Map parameters: {map_params}")
     return create_bpf_map(module, map_name, map_params)
 
 
 def process_bpf_map(func_node, module):
     """Process a BPF map (a function decorated with @map)"""
     map_name = func_node.name
-    print(f"Processing BPF map: {map_name}")
+    logger.info(f"Processing BPF map: {map_name}")
-
-    BPF_MAP_TYPES = {"HashMap": process_hash_map,            # BPF_MAP_TYPE_HASH
-                     "PerfEventArray": process_perf_event_map,   # BPF_MAP_TYPE_PERF_EVENT_ARRAY
-                     }
 
     # For now, assume single return statement
     return_stmt = None
@@ -250,13 +251,13 @@ def process_bpf_map(func_node, module):
 
     rval = return_stmt.value
 
-    # Handle only HashMap maps
     if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name):
-        if rval.func.id in BPF_MAP_TYPES:
+        handler = MapProcessorRegistry.get_processor(rval.func.id)
-            handler = BPF_MAP_TYPES[rval.func.id]
+        if handler:
-            handler(map_name, rval, module)
+            return handler(map_name, rval, module)
         else:
-            print(f"Unknown map type {rval.func.id}, defaulting to HashMap")
+            logger.warning(f"Unknown map type "
-            process_hash_map(map_name, rval, module)
+                           f"{rval.func.id}, defaulting to HashMap")
+            return process_hash_map(map_name, rval, module)
     else:
         raise ValueError("Function under @map must return a map")

pythonbpf/maps/maps_utils.py

@@ -0,0 +1,16 @@
+class MapProcessorRegistry:
+    """Registry for map processor functions"""
+    _processors = {}
+
+    @classmethod
+    def register(cls, map_type_name):
+        """Decorator to register a processor function for a map type"""
+        def decorator(func):
+            cls._processors[map_type_name] = func
+            return func
+        return decorator
+
+    @classmethod
+    def get_processor(cls, map_type_name):
+        """Get the processor function for a map type"""
+        return cls._processors.get(map_type_name)