diff --git a/cocotbext/axi/axis.py b/cocotbext/axi/axis.py index a7b2cd4..e746c1d 100644 --- a/cocotbext/axi/axis.py +++ b/cocotbext/axi/axis.py @@ -552,8 +552,7 @@ class AxiStreamMonitor(AxiStreamBase): self.read_queue = [] - async def recv(self, compact=True): - frame = await self.queue.get() + def _recv(self, frame, compact=True): if self.queue.empty(): self.active_event.clear() self.queue_occupancy_bytes -= len(frame) @@ -562,17 +561,13 @@ class AxiStreamMonitor(AxiStreamBase): frame.compact() return frame + async def recv(self, compact=True): + frame = await self.queue.get() + return self._recv(frame, compact) + def recv_nowait(self, compact=True): - if not self.queue.empty(): - frame = self.queue.get_nowait() - if self.queue.empty(): - self.active_event.clear() - self.queue_occupancy_bytes -= len(frame) - self.queue_occupancy_frames -= 1 - if compact: - frame.compact() - return frame - return None + frame = self.queue.get_nowait() + return self._recv(frame, compact) async def read(self, count=-1): while not self.read_queue: diff --git a/cocotbext/axi/stream.py b/cocotbext/axi/stream.py index d2f4ad8..299172a 100644 --- a/cocotbext/axi/stream.py +++ b/cocotbext/axi/stream.py @@ -241,19 +241,18 @@ class StreamMonitor(StreamBase): _valid_init = None _ready_init = None - async def recv(self): - item = await self.queue.get() + def _recv(self, item): if self.queue.empty(): self.active_event.clear() return item + async def recv(self): + item = await self.queue.get() + return self._recv(item) + def recv_nowait(self): - if not self.queue.empty(): - item = self.queue.get_nowait() - if self.queue.empty(): - self.active_event.clear() - return item - return None + item = self.queue.get_nowait() + return self._recv(item) async def wait(self, timeout=0, timeout_unit=None): if not self.empty():