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>

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()

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):