fix type check and formatting

Author: sohankshirsagar

drift/instrumentation/httpx/e2e-tests/src/app.py

@@ -418,11 +418,13 @@ def test_follow_redirects():
         with httpx.Client(follow_redirects=True) as client:
             # httpbin.org/redirect/2 will redirect twice before returning
             response = client.get("https://httpbin.org/redirect/2")
-            return jsonify({
-                "final_url": str(response.url),
-                "status_code": response.status_code,
-                "redirect_count": len(response.history),
-            })
+            return jsonify(
+                {
+                    "final_url": str(response.url),
+                    "status_code": response.status_code,
+                    "redirect_count": len(response.history),
+                }
+            )
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 

drift/instrumentation/httpx/instrumentation.py

@@ -37,7 +37,7 @@ def __init__(self, message: str, method: str, url: str):
 from ...core.json_schema_helper import DecodedType, EncodingType, SchemaMerge
 from ...core.mode_utils import handle_record_mode, handle_replay_mode
 from ...core.tracing import TdSpanAttributes
-from ...core.tracing.span_utils import CreateSpanOptions, SpanUtils
+from ...core.tracing.span_utils import CreateSpanOptions, SpanInfo, SpanUtils
 from ...core.types import (
     PackageType,
     SpanKind,
@@ -181,21 +181,15 @@ def patched_send(
 
             # Pass through if SDK is disabled
             if sdk.mode == TuskDriftMode.DISABLED:
-                return original_send(
-                    client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
-                )
+                return original_send(client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects)
 
             def original_call():
-                return original_send(
-                    client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
-                )
+                return original_send(client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects)
 
             # REPLAY mode: Use handle_replay_mode for proper background request handling
             if sdk.mode == TuskDriftMode.REPLAY:
                 return handle_replay_mode(
-                    replay_mode_handler=lambda: instrumentation_self._handle_replay_send_sync(
-                        sdk, module, request
-                    ),
+                    replay_mode_handler=lambda: instrumentation_self._handle_replay_send_sync(sdk, module, request),
                     no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
                     is_server_request=False,
                 )
@@ -204,8 +198,13 @@ def original_call():
             return handle_record_mode(
                 original_function_call=original_call,
                 record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_sync(
-                    client_self, request, stream, is_pre_app_start, original_send,
-                    auth=auth, follow_redirects=follow_redirects
+                    client_self,
+                    request,
+                    stream,
+                    is_pre_app_start,
+                    original_send,
+                    auth=auth,
+                    follow_redirects=follow_redirects,
                 ),
                 span_kind=OTelSpanKind.CLIENT,
             )
@@ -277,9 +276,7 @@ def _handle_record_send_sync(
             with SpanUtils.with_span(span_info):
                 # Check drop transforms BEFORE making the request
                 headers = dict(request.headers)
-                if self._transform_engine and self._transform_engine.should_drop_outbound_request(
-                    method, url, headers
-                ):
+                if self._transform_engine and self._transform_engine.should_drop_outbound_request(method, url, headers):
                     # Request should be dropped - mark span and raise exception
                     span_info.span.set_attribute(
                         TdSpanAttributes.OUTPUT_VALUE,
@@ -357,9 +354,7 @@ async def original_call():
             # handle_replay_mode returns coroutine which we await
             if sdk.mode == TuskDriftMode.REPLAY:
                 return await handle_replay_mode(
-                    replay_mode_handler=lambda: instrumentation_self._handle_replay_send_async(
-                        sdk, module, request
-                    ),
+                    replay_mode_handler=lambda: instrumentation_self._handle_replay_send_async(sdk, module, request),
                     no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
                     is_server_request=False,
                 )
@@ -369,8 +364,13 @@ async def original_call():
             return await handle_record_mode(
                 original_function_call=original_call,
                 record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_async(
-                    client_self, request, stream, is_pre_app_start, original_send,
-                    auth=auth, follow_redirects=follow_redirects
+                    client_self,
+                    request,
+                    stream,
+                    is_pre_app_start,
+                    original_send,
+                    auth=auth,
+                    follow_redirects=follow_redirects,
                 ),
                 span_kind=OTelSpanKind.CLIENT,
             )
@@ -418,9 +418,7 @@ async def _handle_record_send_async(
             with SpanUtils.with_span(span_info):
                 # Check drop transforms BEFORE making the request
                 headers = dict(request.headers)
-                if self._transform_engine and self._transform_engine.should_drop_outbound_request(
-                    method, url, headers
-                ):
+                if self._transform_engine and self._transform_engine.should_drop_outbound_request(method, url, headers):
                     # Request should be dropped - mark span and raise exception
                     span_info.span.set_attribute(
                         TdSpanAttributes.OUTPUT_VALUE,
@@ -538,7 +536,7 @@ def _get_request_body_safely(self, request: Any) -> bytes | None:
         """
         try:
             # Check if content is already available (non-streaming requests)
-            if hasattr(request, '_content') and request._content is not None:
+            if hasattr(request, "_content") and request._content is not None:
                 return request.content
 
             # For streaming/multipart requests, read the content
@@ -578,7 +576,7 @@ def _try_get_mock_from_request_sync(
             # Extract query params from URL
             # httpx.URL has a .params attribute that returns QueryParams
             params = {}
-            if hasattr(request.url, 'params'):
+            if hasattr(request.url, "params"):
                 params = dict(request.url.params)
 
             # Get body from request - handle streaming/multipart bodies
@@ -701,9 +699,7 @@ def _create_mock_response(self, httpx_module: Any, mock_data: dict[str, Any], me
             # Create minimal placeholder Response objects for history
             # These represent the intermediate redirect responses
             for i in range(history_count):
-                redirect_request = httpx_module.Request(
-                    method.upper(), url if i == 0 else f"redirect_{i}"
-                )
+                redirect_request = httpx_module.Request(method.upper(), url if i == 0 else f"redirect_{i}")
                 redirect_response = httpx_module.Response(
                     status_code=302,  # Standard redirect status
                     content=b"",
@@ -753,7 +749,7 @@ def _finalize_span_from_request(
 
             # Extract query params from URL
             params = {}
-            if hasattr(request.url, 'params'):
+            if hasattr(request.url, "params"):
                 params = dict(request.url.params)
 
             # Get request body - handle streaming/multipart bodies