Add async callbacks support for apply_on_packets

src/pyshark/capture/capture.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import os
 import threading
 import subprocess
@@ -291,7 +292,10 @@ async def _go_through_packets_from_fd(self, fd, packet_callback, packet_count=No
             if packet:
                 packets_captured += 1
                 try:
-                    packet_callback(packet)
+                    if inspect.iscoroutinefunction(packet_callback):
+                        await packet_callback(packet)
+                    else:
+                        packet_callback(packet)
                 except StopCapture:
                     self._log.debug("User-initiated capture stop in callback")
                     break

src/setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="pyshark",
-    version="0.6",
+    version="0.6.1",
     packages=find_packages(),
     package_data={'': ['*.ini', '*.pcapng']},
     install_requires=['lxml', 'termcolor', 'packaging', 'appdirs'],

tests/test_cap_operations.py

@@ -17,6 +17,14 @@ def test_packet_callback_called_for_each_packet(lazy_simple_capture):
     assert mock_callback.call_count == 24
 
 
+def test_async_packet_callback_called_for_each_packet(lazy_simple_capture):
+    # Test cap has 24 packets
+    mock_callback = mock.AsyncMock()
+    lazy_simple_capture.apply_on_packets(mock_callback)
+    assert mock_callback.call_count == 24
+    mock_callback.assert_awaited()
+
+
 def test_apply_on_packet_stops_on_timeout(lazy_simple_capture):
     def wait(pkt):
         time.sleep(5)