Add support for custom aiohttp connector (#325)

This should allow users that need it to add support for SOCKS5 in their application via https://pypi.org/project/aiohttp-socks/. Partially fixes: #147 Signed-off-by: rany <[email protected]>
rany2 · Nov 22, 2024 · d619925 · d619925
1 parent 7ee40ff
commit d619925
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git a/src/edge_tts/communicate.py b/src/edge_tts/communicate.py
@@ -238,7 +238,7 @@ def calc_max_mesg_size(tts_config: TTSConfig) -> int:
 
 class Communicate:
     """
-    Class for communicating with the service.
+    Communicate with the service.
     """
 
     def __init__(
@@ -249,17 +249,11 @@ def __init__(
         rate: str = "+0%",
         volume: str = "+0%",
         pitch: str = "+0Hz",
+        connector: Optional[aiohttp.BaseConnector] = None,
         proxy: Optional[str] = None,
-        connect_timeout: int = 10,
-        receive_timeout: int = 60,
+        connect_timeout: Optional[int] = 10,
+        receive_timeout: Optional[int] = 60,
     ):
-        """
-        Initializes the Communicate class.
-
-        Raises:
-            ValueError: If the voice is not valid.
-        """
-
         # Validate TTS settings and store the TTSConfig object.
         self.tts_config = TTSConfig(voice, rate, volume, pitch)
 
@@ -290,6 +284,11 @@ def __init__(
             sock_read=receive_timeout,
         )
 
+        # Validate the connector parameter.
+        if connector is not None and not isinstance(connector, aiohttp.BaseConnector):
+            raise TypeError("connector must be aiohttp.BaseConnector")
+        self.connector: Optional[aiohttp.BaseConnector] = connector
+
         # Store current state of TTS.
         self.state: Dict[str, Any] = {
             "partial_text": None,
@@ -351,6 +350,7 @@ async def send_ssml_request() -> None:
         # Create a new connection to the service.
         ssl_ctx = ssl.create_default_context(cafile=certifi.where())
         async with aiohttp.ClientSession(
+            connector=self.connector,
             trust_env=True,
             timeout=self.session_timeout,
         ) as session, session.ws_connect(

diff --git a/src/edge_tts/submaker.py b/src/edge_tts/submaker.py
@@ -15,7 +15,7 @@ def __init__(self) -> None:
 
     def add_cue(self, timestamp: Tuple[float, float], text: str) -> None:
         """
-        Add a subtitle part to the SubMaker object.
+        Add a cue to the SubMaker object.
 
         Args:
             timestamp (tuple): The offset and duration of the subtitle.

diff --git a/src/edge_tts/voices.py b/src/edge_tts/voices.py
@@ -55,21 +55,24 @@ async def __list_voices(
     return data
 
 
-async def list_voices(*, proxy: Optional[str] = None) -> List[Voice]:
+async def list_voices(
+    *, connector: Optional[aiohttp.BaseConnector] = None, proxy: Optional[str] = None
+) -> List[Voice]:
     """
     List all available voices and their attributes.
 
     This pulls data from the URL used by Microsoft Edge to return a list of
     all available voices.
 
     Args:
+        connector (Optional[aiohttp.BaseConnector]): The connector to use for the request.
         proxy (Optional[str]): The proxy to use for the request.
 
     Returns:
         List[Voice]: A list of voices and their attributes.
     """
     ssl_ctx = ssl.create_default_context(cafile=certifi.where())
-    async with aiohttp.ClientSession(trust_env=True) as session:
+    async with aiohttp.ClientSession(connector=connector, trust_env=True) as session:
         try:
             data = await __list_voices(session, ssl_ctx, proxy)
         except aiohttp.ClientResponseError as e: