Merge setup functions (#277)

* Merge setup functions

* Add to changelog

CHANGELOG.md

@@ -5,6 +5,7 @@
 * Deprecation of MetricsProvider interface. For the backward compatibility, prometheus-client is conditionally imported. To use it, install prometheus-client. Related PRs: [#271](https://github.com/anna-money/aio-request/pull/271), [#218](https://github.com/anna-money/aio-request/pull/218), [#268](https://github.com/anna-money/aio-request/pull/268)
 * [Removal of unused Client interface](https://github.com/anna-money/aio-request/commit/fe75660af8e7520a6fa5143f982c5aacd2ea079a)
 * [Do not retry low timeout response](https://github.com/anna-money/aio-request/pull/276)
+* [Merge setup functions](https://github.com/anna-money/aio-request/pull/277)
 
 
 ## v0.1.32 (2024-10-18)

aio_request/__init__.py

@@ -21,7 +21,18 @@
 from .endpoint_provider import DelegateEndpointProvider, EndpointProvider, StaticEndpointProvider
 from .pipeline import BypassModule, LowTimeoutModule, NextModuleFunc, RequestModule, TransportModule, build_pipeline
 from .priority import Priority
-from .request import delete, get, patch, patch_json, post, post_json, put, put_json
+from .request import (
+    delete,
+    get,
+    patch,
+    patch_json,
+    post,
+    post_json,
+    put,
+    put_json,
+    RequestEnricher,
+    AsyncRequestEnricher,
+)
 from .request_strategy import (
     MethodBasedStrategy,
     ParallelRequestStrategy,
@@ -94,6 +105,8 @@
     "post_json",
     "put",
     "put_json",
+    "RequestEnricher",
+    "AsyncRequestEnricher",
     # request_strategy.py
     "MethodBasedStrategy",
     "ParallelRequestStrategy",

aio_request/pipeline.py

@@ -1,4 +1,5 @@
 import abc
+import asyncio
 import collections.abc
 
 import multidict
@@ -8,6 +9,7 @@
 from .circuit_breaker import CircuitBreaker
 from .deadline import Deadline
 from .priority import Priority
+from .request import AsyncRequestEnricher, RequestEnricher
 from .response_classifier import ResponseClassifier, ResponseVerdict
 from .transport import Transport
 
@@ -82,7 +84,7 @@ def __init__(
         transport: Transport,
         *,
         emit_system_headers: bool,
-        request_enricher: collections.abc.Callable[[Request, bool], collections.abc.Awaitable[Request]] | None,
+        request_enricher: RequestEnricher | AsyncRequestEnricher | None,
     ):
         self.__transport = transport
         self.__emit_system_headers = emit_system_headers
@@ -106,11 +108,15 @@ async def execute(
                 }
             )
 
-        request = (
-            await self.__request_enricher(request, self.__emit_system_headers)
-            if self.__request_enricher is not None
-            else request
-        )
+        if self.__request_enricher is not None:
+            request = (
+                await enriched_request  # type: ignore
+                if (
+                    (enriched_request := self.__request_enricher(request, self.__emit_system_headers))
+                    and asyncio.iscoroutine(enriched_request)
+                )
+                else enriched_request
+            )
 
         return await self.__transport.send(endpoint, request, deadline.timeout)
 

aio_request/request.py

@@ -7,6 +7,9 @@
 
 from .base import MAX_REDIRECTS, Header, Headers, Method, PathParameters, QueryParameters, Request
 
+RequestEnricher = collections.abc.Callable[[Request, bool], Request]
+AsyncRequestEnricher = collections.abc.Callable[[Request, bool], collections.abc.Awaitable[Request]]
+
 
 def get(
     url: str | yarl.URL,

aio_request/setup.py

@@ -1,60 +1,24 @@
-import collections.abc
 import warnings
 from typing import Any
 
 import yarl
 
-from .base import ClosableResponse, Method, Request
+from .base import ClosableResponse, Method
 from .circuit_breaker import CircuitBreaker
 from .client import Client
 from .delays_provider import linear_backoff_delays
 from .deprecated import MetricsProvider
 from .endpoint_provider import EndpointProvider, StaticEndpointProvider
 from .pipeline import BypassModule, CircuitBreakerModule, LowTimeoutModule, TransportModule, build_pipeline
 from .priority import Priority
+from .request import AsyncRequestEnricher, RequestEnricher
 from .request_strategy import MethodBasedStrategy, RequestStrategy, sequential_strategy, single_attempt_strategy
 from .response_classifier import DefaultResponseClassifier, ResponseClassifier
 from .transport import Transport
 
 MISSING: Any = object()
 
 
-def setup(
-    *,
-    transport: Transport,
-    endpoint: str | yarl.URL = MISSING,
-    endpoint_provider: EndpointProvider = MISSING,
-    safe_method_strategy: RequestStrategy = sequential_strategy(
-        attempts_count=3, delays_provider=linear_backoff_delays()
-    ),
-    unsafe_method_strategy: RequestStrategy = single_attempt_strategy(),
-    response_classifier: ResponseClassifier | None = None,
-    timeout: float = 20.0,
-    priority: Priority = Priority.NORMAL,
-    low_timeout_threshold: float = 0.005,
-    emit_system_headers: bool = True,
-    request_enricher: collections.abc.Callable[[Request], Request] | None = None,
-    circuit_breaker: CircuitBreaker[yarl.URL, ClosableResponse] | None = None,
-) -> Client:
-    async def _enrich_request(request: Request, _: bool) -> Request:
-        return request_enricher(request) if request_enricher is not None else request
-
-    return setup_v2(
-        transport=transport,
-        endpoint=endpoint,
-        endpoint_provider=endpoint_provider,
-        safe_method_strategy=safe_method_strategy,
-        unsafe_method_strategy=unsafe_method_strategy,
-        response_classifier=response_classifier,
-        timeout=timeout,
-        priority=priority,
-        low_timeout_threshold=low_timeout_threshold,
-        emit_system_headers=emit_system_headers,
-        request_enricher=_enrich_request,
-        circuit_breaker=circuit_breaker,
-    )
-
-
 def setup_v2(
     *,
     transport: Transport,
@@ -69,7 +33,7 @@ def setup_v2(
     priority: Priority = Priority.NORMAL,
     low_timeout_threshold: float = 0.005,
     emit_system_headers: bool = True,
-    request_enricher: collections.abc.Callable[[Request, bool], collections.abc.Awaitable[Request]] | None = None,
+    request_enricher: RequestEnricher | AsyncRequestEnricher | None = None,
     metrics_provider: MetricsProvider | None = None,
     circuit_breaker: CircuitBreaker[yarl.URL, ClosableResponse] | None = None,
 ) -> Client:
@@ -118,3 +82,6 @@ def setup_v2(
             ],
         ),
     )
+
+
+setup = setup_v2