chore: Add Type Annotations

Signed-off-by: Varsha GS <varsha.gs@ibm.com>

Author: GSVarsha

src/instana/instrumentation/aio_pika.py

@@ -3,21 +3,44 @@
 try:
     import aio_pika
     import wrapt
+    from typing import (
+        TYPE_CHECKING,
+        Dict,
+        Any,
+        Callable,
+        Tuple,
+        Type,
+        Optional,
+    )
 
     from instana.log import logger
     from instana.propagators.format import Format
     from instana.util.traceutils import get_tracer_tuple, tracing_is_off
     from instana.singletons import tracer
 
-    def _extract_span_attributes(span, connection, sort, routing_key, exchange) -> None:
+    if TYPE_CHECKING:
+        from instana.span.span import InstanaSpan
+        from aio_pika.exchange import Exchange
+        from aiormq.abc import ConfirmationFrameType
+        from aio_pika.abc import ConsumerTag, AbstractMessage
+        from aio_pika.queue import Queue, QueueIterator
+
+    def _extract_span_attributes(
+        span: "InstanaSpan", connection, sort: str, routing_key: str, exchange: str
+    ) -> None:
         span.set_attribute("address", str(connection.url))
 
         span.set_attribute("sort", sort)
         span.set_attribute("key", routing_key)
         span.set_attribute("exchange", exchange)
 
     @wrapt.patch_function_wrapper("aio_pika", "Exchange.publish")
-    async def publish_with_instana(wrapped, instance, args, kwargs):
+    async def publish_with_instana(
+        wrapped: Callable[..., Optional["ConfirmationFrameType"]],
+        instance: "Exchange",
+        args: Tuple[object],
+        kwargs: Dict[str, Any],
+    ) -> Optional["ConfirmationFrameType"]:
         if tracing_is_off():
             return await wrapped(*args, **kwargs)
 
@@ -47,20 +70,32 @@ async def publish_with_instana(wrapped, instance, args, kwargs):
                 return response
 
     @wrapt.patch_function_wrapper("aio_pika", "Queue.consume")
-    async def consume_with_instana(wrapped, instance, args, kwargs):
+    async def consume_with_instana(
+        wrapped: Callable[..., "ConsumerTag"],
+        instance: Type["Queue"],
+        args: Tuple[object],
+        kwargs: Dict[str, Any],
+    ) -> "ConsumerTag":
         connection = instance.channel._connection
         callback = kwargs["callback"] if kwargs.get("callback") else args[0]
 
         @wrapt.decorator
-        async def callback_wrapper(wrapped, instance, args, kwargs):
+        async def callback_wrapper(
+            wrapped: Callable[[Type["AbstractMessage"]], Any],
+            instance: Type["QueueIterator"],
+            args: Tuple[Type["AbstractMessage"], ...],
+            kwargs: Dict[str, Any],
+        ) -> Callable[[Type["AbstractMessage"]], Any]:
             message = args[0]
             parent_context = tracer.extract(
                 Format.HTTP_HEADERS, message.headers, disable_w3c_trace_context=True
             )
             with tracer.start_as_current_span(
                 "rabbitmq", span_context=parent_context
             ) as span:
-                _extract_span_attributes(span, connection, "consume", message.routing_key, message.exchange)
+                _extract_span_attributes(
+                    span, connection, "consume", message.routing_key, message.exchange
+                )
             try:
                 response = await wrapped(*args, **kwargs)
             except Exception as exc:
@@ -75,7 +110,7 @@ async def callback_wrapper(wrapped, instance, args, kwargs):
             args = (wrapped_callback,) + args[1:]
 
         return await wrapped(*args, **kwargs)
-    
+
     logger.debug("Instrumenting aio-pika")
 
 except ImportError:

tests/clients/test_aio_pika.py

@@ -1,12 +1,15 @@
 # (c) Copyright IBM Corp. 2025
 
 import pytest
-from typing import Generator
+from typing import Generator, TYPE_CHECKING
 import asyncio
 from aio_pika import Message, connect, connect_robust
 
 from instana.singletons import agent, tracer
 
+if TYPE_CHECKING:
+    from instana.span.readable_span import ReadableSpan
+
 
 class TestAioPika:
     @pytest.fixture(autouse=True)
@@ -15,7 +18,7 @@ def _resource(self) -> Generator[None, None, None]:
         # setup
         self.recorder = tracer.span_processor
         self.recorder.clear_spans()
-        
+
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(None)
         self.queue_name = "test.queue"
@@ -155,7 +158,7 @@ def test_basic_consume(self, connect_method) -> None:
         assert not test_span.ec
 
         # Span attributes
-        def assert_span_info(rabbitmq_span, sort) -> None:
+        def assert_span_info(rabbitmq_span: "ReadableSpan", sort: str) -> None:
             assert rabbitmq_span.data["rabbitmq"]["exchange"] == "test.exchange"
             assert rabbitmq_span.data["rabbitmq"]["sort"] == sort
             assert rabbitmq_span.data["rabbitmq"]["address"]