Add __contains__ for Function, Section, Segment, BasicBlock and BinaryView

Addresses issue binaryninja-api#1113

3 files changed, 20 insertions, 0 deletions

diff --git a/python/basicblock.py b/python/basicblock.py
index 623cc3ed..872e44b8 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -129,6 +129,9 @@ class BasicBlock:
 		data = self.view.read(start, length)
 		return self.arch.get_instruction_text(data, start)
 
+	def __contains__(self, i:int):
+		return i >= self.start and i < self.end
+
 	def _buildStartCache(self) -> None:
 		if self._instStarts is None:
 			# build the instruction start cache
diff --git a/python/binaryview.py b/python/binaryview.py
index 12f886fd..ef559d8b 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -978,6 +978,9 @@ class Segment:
 	def __hash__(self):
 		return hash(ctypes.addressof(self.handle.contents))
 
+	def __contains__(self, i:int):
+		return i >= self.start and i < self.end
+
 	@property
 	def start(self) -> int:
 		return core.BNSegmentGetStart(self.handle)
@@ -1073,6 +1076,9 @@ class Section:
 	def __hash__(self):
 		return hash(ctypes.addressof(self.handle.contents))
 
+	def __contains__(self, i:int):
+		return i >= self.start and i < self.end
+
 	@property
 	def name(self) -> str:
 		return core.BNSectionGetName(self.handle)
@@ -1493,6 +1499,12 @@ class BinaryView:
 		else:
 			raise IndexError("index out of range")
 
+	def __contains__(self, i:int):
+		for s in self.segments:
+			if i in s:
+				return True
+		return False
+
 	@classmethod
 	def register(cls) -> None:
 		binaryninja._init_plugins()
diff --git a/python/function.py b/python/function.py
index 42e88fa3..cedfb9e7 100644
--- a/python/function.py
+++ b/python/function.py
@@ -264,6 +264,11 @@ class Function:
 			result += token.text
 		return result
 
+	def __contains__(self, value:Union[basicblock.BasicBlock, int]):
+		if isinstance(value, basicblock.BasicBlock):
+			return value.function == self
+		return self in [block.function for block in self.view.get_basic_blocks_at(int(value))]
+
 	@classmethod
 	def _unregister(cls, func:'core.BNFunction') -> None:
 		handle = ctypes.cast(func, ctypes.c_void_p)