Skip to content

Commit cb32201

Browse files
committed
Add Address
1 parent b61cd6f commit cb32201

File tree

1 file changed

+35
-28
lines changed

1 file changed

+35
-28
lines changed

vm/src/stdlib/socket.rs

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,9 @@ use std::io::Read;
44
use std::io::Write;
55
use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket};
66

7-
use crate::function::PyFuncArgs;
87
use crate::obj::objbytes::PyBytesRef;
9-
use crate::obj::objint;
108
use crate::obj::objint::PyIntRef;
11-
use crate::obj::objstr;
9+
use crate::obj::objstr::PyStringRef;
1210
use crate::obj::objtuple::PyTupleRef;
1311
use crate::pyobject::{PyObjectRef, PyRef, PyResult, PyValue, TryFromObject};
1412
use crate::vm::VirtualMachine;
@@ -173,8 +171,8 @@ impl SocketRef {
173171
Socket::new(family, kind).into_ref_with_type(vm, cls)
174172
}
175173

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();
178176

179177
match self.socket_kind {
180178
SocketKind::Stream => match TcpStream::connect(address_string) {
@@ -197,8 +195,8 @@ impl SocketRef {
197195
}
198196
}
199197

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();
202200

203201
match self.socket_kind {
204202
SocketKind::Stream => match TcpListener::bind(address_string) {
@@ -285,8 +283,8 @@ impl SocketRef {
285283
Ok(())
286284
}
287285

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();
290288

291289
match self.socket_kind {
292290
SocketKind::Dgram => {
@@ -337,25 +335,34 @@ impl SocketRef {
337335
}
338336
}
339337

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+
}
359366
}
360367

361368
fn get_addr_tuple(vm: &VirtualMachine, addr: SocketAddr) -> PyResult {

0 commit comments

Comments
 (0)