From 2483ef28404681ab1117e983194c06ffa24a1a16 Mon Sep 17 00:00:00 2001
From: varun-r-mallya <varunrmallya@gmail.com>
Date: Sat, 11 Oct 2025 04:19:26 +0530
Subject: [PATCH] separate vmlinux class handler

---
 .gitignore                                    |  1 +
 pythonbpf/vmlinux_parser/dependency_data.py   |  1 +
 pythonbpf/vmlinux_parser/import_detector.py   | 27 ++++++++++++++++---
 .../vmlinux_parser/vmlinux_class_handler.py   |  8 ++++++
 4 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 pythonbpf/vmlinux_parser/dependency_data.py
 create mode 100644 pythonbpf/vmlinux_parser/vmlinux_class_handler.py

diff --git a/.gitignore b/.gitignore
index 085331b..585a95e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@ __pycache__/
 *.o
 .ipynb_checkpoints/
 vmlinux.py
+~*
diff --git a/pythonbpf/vmlinux_parser/dependency_data.py b/pythonbpf/vmlinux_parser/dependency_data.py
new file mode 100644
index 0000000..0b9ee51
--- /dev/null
+++ b/pythonbpf/vmlinux_parser/dependency_data.py
@@ -0,0 +1 @@
+# This will contain a node class that has special properties to be a dependency tree node
diff --git a/pythonbpf/vmlinux_parser/import_detector.py b/pythonbpf/vmlinux_parser/import_detector.py
index 1cb3808..e9be65c 100644
--- a/pythonbpf/vmlinux_parser/import_detector.py
+++ b/pythonbpf/vmlinux_parser/import_detector.py
@@ -1,7 +1,7 @@
 import ast
 import logging
-import llvmlite.ir as ir
 from typing import List, Tuple
+from .vmlinux_class_handler import process_vmlinux_class
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +70,27 @@ def detect_import_statement(tree: ast.AST) -> List[Tuple[str, str]]:
     logger.info(f"Total vmlinux imports detected: {len(vmlinux_imports)}")
     return vmlinux_imports
 
-def vmlinux_proc(tree, module):
+def vmlinux_proc(tree: ast.AST, module):
     import_statements = detect_import_statement(tree)
-    logger.info(f"Import statements {import_statements}")
+
+    if not import_statements:
+        logger.info("No vmlinux imports found")
+        return
+
+    vmlinux_types = set()
+    for module_name, imported_item in import_statements:
+        vmlinux_types.add(imported_item)
+        logger.info(f"Registered vmlinux type: {imported_item}")
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef):
+            # Check if this class uses vmlinux types
+            logger.info(f"Processing ClassDef with vmlinux types: {node.name}")
+            process_vmlinux_class(node, module, vmlinux_types)
+
+        elif isinstance(node, ast.Assign):
+            logger.info(f"Processing Assign with vmlinux types")
+            process_vmlinux_assign(node, module, vmlinux_types)
+
+def process_vmlinux_assign(node, module, vmlinux_types):
+    raise NotImplementedError("Assignment handling has not been implemented yet")
diff --git a/pythonbpf/vmlinux_parser/vmlinux_class_handler.py b/pythonbpf/vmlinux_parser/vmlinux_class_handler.py
new file mode 100644
index 0000000..94787dd
--- /dev/null
+++ b/pythonbpf/vmlinux_parser/vmlinux_class_handler.py
@@ -0,0 +1,8 @@
+import ast
+import logging
+
+logger = logging.getLogger(__name__)
+
+def process_vmlinux_class(node, module, vmlinux_types):
+    # Process ClassDef nodes that use vmlinux imports
+    pass