@@ -4,11 +4,9 @@ use std::io::Read;
4
4
use std:: io:: Write ;
5
5
use std:: net:: { SocketAddr , TcpListener , TcpStream , ToSocketAddrs , UdpSocket } ;
6
6
7
- use crate :: function:: PyFuncArgs ;
8
7
use crate :: obj:: objbytes:: PyBytesRef ;
9
- use crate :: obj:: objint;
10
8
use crate :: obj:: objint:: PyIntRef ;
11
- use crate :: obj:: objstr;
9
+ use crate :: obj:: objstr:: PyStringRef ;
12
10
use crate :: obj:: objtuple:: PyTupleRef ;
13
11
use crate :: pyobject:: { PyObjectRef , PyRef , PyResult , PyValue , TryFromObject } ;
14
12
use crate :: vm:: VirtualMachine ;
@@ -173,8 +171,8 @@ impl SocketRef {
173
171
Socket :: new ( family, kind) . into_ref_with_type ( vm, cls)
174
172
}
175
173
176
- fn connect ( self , address : PyTupleRef , vm : & VirtualMachine ) -> PyResult < ( ) > {
177
- let address_string = get_address_string ( vm , address ) ? ;
174
+ fn connect ( self , address : Address , vm : & VirtualMachine ) -> PyResult < ( ) > {
175
+ let address_string = address . get_address_string ( ) ;
178
176
179
177
match self . socket_kind {
180
178
SocketKind :: Stream => match TcpStream :: connect ( address_string) {
@@ -197,8 +195,8 @@ impl SocketRef {
197
195
}
198
196
}
199
197
200
- fn bind ( self , address : PyTupleRef , vm : & VirtualMachine ) -> PyResult < ( ) > {
201
- let address_string = get_address_string ( vm , address ) ? ;
198
+ fn bind ( self , address : Address , vm : & VirtualMachine ) -> PyResult < ( ) > {
199
+ let address_string = address . get_address_string ( ) ;
202
200
203
201
match self . socket_kind {
204
202
SocketKind :: Stream => match TcpListener :: bind ( address_string) {
@@ -285,8 +283,8 @@ impl SocketRef {
285
283
Ok ( ( ) )
286
284
}
287
285
288
- fn sendto ( self , bytes : PyBytesRef , address : PyTupleRef , vm : & VirtualMachine ) -> PyResult < ( ) > {
289
- let address_string = get_address_string ( vm , address ) ? ;
286
+ fn sendto ( self , bytes : PyBytesRef , address : Address , vm : & VirtualMachine ) -> PyResult < ( ) > {
287
+ let address_string = address . get_address_string ( ) ;
290
288
291
289
match self . socket_kind {
292
290
SocketKind :: Dgram => {
@@ -337,25 +335,34 @@ impl SocketRef {
337
335
}
338
336
}
339
337
340
- fn get_address_string ( vm : & VirtualMachine , address : PyTupleRef ) -> Result < String , PyObjectRef > {
341
- let args = PyFuncArgs {
342
- args : address. elements . borrow ( ) . to_vec ( ) ,
343
- kwargs : vec ! [ ] ,
344
- } ;
345
- arg_check ! (
346
- vm,
347
- args,
348
- required = [
349
- ( host, Some ( vm. ctx. str_type( ) ) ) ,
350
- ( port, Some ( vm. ctx. int_type( ) ) )
351
- ]
352
- ) ;
353
-
354
- Ok ( format ! (
355
- "{}:{}" ,
356
- objstr:: get_value( host) ,
357
- objint:: get_value( port) . to_string( )
358
- ) )
338
+ struct Address {
339
+ host : String ,
340
+ port : usize ,
341
+ }
342
+
343
+ impl Address {
344
+ fn get_address_string ( self ) -> String {
345
+ format ! ( "{}:{}" , self . host, self . port. to_string( ) )
346
+ }
347
+ }
348
+
349
+ impl TryFromObject for Address {
350
+ fn try_from_object ( vm : & VirtualMachine , obj : PyObjectRef ) -> PyResult < Self > {
351
+ let tuple = PyTupleRef :: try_from_object ( vm, obj) ?;
352
+ if tuple. elements . borrow ( ) . len ( ) != 2 {
353
+ Err ( vm. new_type_error ( "Address tuple should have only 2 values" . to_string ( ) ) )
354
+ } else {
355
+ Ok ( Address {
356
+ host : PyStringRef :: try_from_object ( vm, tuple. elements . borrow ( ) [ 0 ] . clone ( ) ) ?
357
+ . value
358
+ . to_string ( ) ,
359
+ port : PyIntRef :: try_from_object ( vm, tuple. elements . borrow ( ) [ 1 ] . clone ( ) ) ?
360
+ . as_bigint ( )
361
+ . to_usize ( )
362
+ . unwrap ( ) ,
363
+ } )
364
+ }
365
+ }
359
366
}
360
367
361
368
fn get_addr_tuple ( vm : & VirtualMachine , addr : SocketAddr ) -> PyResult {
0 commit comments