Add getitem api to BPF object to reference tables

To simplify some common usage patterns, add a getitem api to the BPF
class. This wraps the get_table() api and allows the api user to skip
keeping their own get_table handle.

Fixes: #137
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

diff --git a/examples/vfsreadlat.py b/examples/vfsreadlat.py
index da6a0d4..febd5c4 100755 (executable)
--- a/examples/vfsreadlat.py
+++ b/examples/vfsreadlat.py
@@ -41,7 +41,6 @@ if len(argv) > 1:
 b = BPF(src_file = "vfsreadlat.c")
 b.attach_kprobe(event="vfs_read", fn_name="do_entry")
 b.attach_kretprobe(event="vfs_read", fn_name="do_return")
-dist = b.get_table("dist")
 dist_max = 64
 
 # header
@@ -69,7 +68,7 @@ def print_log2_hist(d, val_type):
        val_max = 0
        for i in range(1, dist_max + 1):
                try:
-                       val = dist[c_int(i)].value - last[i]
+                       val = b["dist"][c_int(i)].value - last[i]
                        if (val > 0):
                                idx_max = i
                        if (val > val_max):
@@ -84,10 +83,10 @@ def print_log2_hist(d, val_type):
                if (low == high):
                        low -= 1
                try:
-                       val = dist[c_int(i)].value - last[i]
+                       val = b["dist"][c_int(i)].value - last[i]
                        print("%8d -> %-8d : %-8d |%-*s|" % (low, high, val,
                            stars_max, stars(val, val_max, stars_max)))
-                       last[i] = dist[c_int(i)].value
+                       last[i] = b["dist"][c_int(i)].value
                except:
                        break
 
@@ -105,6 +104,6 @@ while (1):
                pass; do_exit = 1
 
        print
-       print_log2_hist(dist, "usecs")
+       print_log2_hist(b["dist"], "usecs")
        if do_exit:
                exit()

diff --git a/src/python/bpf/__init__.py b/src/python/bpf/__init__.py
index ae74467..3233b07 100644 (file)
--- a/src/python/bpf/__init__.py
+++ b/src/python/bpf/__init__.py
@@ -223,6 +223,7 @@ class BPF(object):
     def __init__(self, src_file="", hdr_file="", text=None, debug=0):
         self.debug = debug
         self.funcs = {}
+        self.tables = {}
         if text:
             self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"), self.debug)
         else:
@@ -297,7 +298,7 @@ class BPF(object):
         map_id = lib.bpf_table_id(self.module, name.encode("ascii"))
         map_fd = lib.bpf_table_fd(self.module, name.encode("ascii"))
         if map_fd < 0:
-            raise Exception("Failed to find BPF Table %s" % name)
+            raise KeyError
         if not keytype:
             key_desc = lib.bpf_table_key_desc(self.module, name.encode("ascii"))
             if not key_desc:
@@ -310,6 +311,23 @@ class BPF(object):
             leaftype = BPF._decode_table_type(json.loads(leaf_desc.decode()))
         return BPF.Table(self, map_id, map_fd, keytype, leaftype)
 
+    def __getitem__(self, key):
+        if key not in self.tables:
+            self.tables[key] = self.get_table(key)
+        return self.tables[key]
+
+    def __setitem__(self, key, leaf):
+        self.tables[key] = leaf
+
+    def __len__(self):
+        return len(self.tables)
+
+    def __delitem__(self, key):
+        del self.tables[key]
+
+    def __iter__(self):
+        return self.tables.__iter__()
+
     @staticmethod
     def attach_raw_socket(fn, dev):
         if not isinstance(fn, BPF.Function):