@@ -178,33 +178,30 @@ async def open_tcp_stream(
178178 ssl_context : typing .Optional [ssl .SSLContext ],
179179 timeout : TimeoutConfig ,
180180 ) -> SocketStream :
181- return await self ._open_stream (
182- trio .open_tcp_stream (hostname , port ), hostname , ssl_context , timeout
183- )
181+ connect_timeout = _or_inf (timeout .connect_timeout )
182+
183+ with trio .move_on_after (connect_timeout ) as cancel_scope :
184+ stream : trio .SocketStream = await trio .open_tcp_stream (hostname , port )
185+ if ssl_context is not None :
186+ stream = trio .SSLStream (stream , ssl_context , server_hostname = hostname )
187+ await stream .do_handshake ()
188+
189+ if cancel_scope .cancelled_caught :
190+ raise ConnectTimeout ()
191+
192+ return SocketStream (stream = stream , timeout = timeout )
184193
185194 async def open_uds_stream (
186195 self ,
187196 path : str ,
188197 hostname : typing .Optional [str ],
189198 ssl_context : typing .Optional [ssl .SSLContext ],
190199 timeout : TimeoutConfig ,
191- ) -> SocketStream :
192- hostname = hostname if ssl_context else None
193- return await self ._open_stream (
194- trio .open_unix_socket (path ), hostname , ssl_context , timeout
195- )
196-
197- async def _open_stream (
198- self ,
199- socket_stream : typing .Awaitable [trio .SocketStream ],
200- hostname : typing .Optional [str ],
201- ssl_context : typing .Optional [ssl .SSLContext ],
202- timeout : TimeoutConfig ,
203200 ) -> SocketStream :
204201 connect_timeout = _or_inf (timeout .connect_timeout )
205202
206203 with trio .move_on_after (connect_timeout ) as cancel_scope :
207- stream : trio .SocketStream = await socket_stream
204+ stream : trio .SocketStream = await trio . open_unix_socket ( path )
208205 if ssl_context is not None :
209206 stream = trio .SSLStream (stream , ssl_context , server_hostname = hostname )
210207 await stream .do_handshake ()
0 commit comments