Skip to content

Commit 59d4e4e

Browse files
committed
Fix os_open
1 parent 0280a1c commit 59d4e4e

File tree

3 files changed

+87
-78
lines changed

3 files changed

+87
-78
lines changed

vm/src/pyobject.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1069,6 +1069,15 @@ pub enum Either<A, B> {
10691069
B(B),
10701070
}
10711071

1072+
impl<A: PyValue, B: PyValue> Either<PyRef<A>, PyRef<B>> {
1073+
pub fn into_object(self) -> PyObjectRef {
1074+
match self {
1075+
Either::A(a) => a.into_object(),
1076+
Either::B(b) => b.into_object(),
1077+
}
1078+
}
1079+
}
1080+
10721081
/// This allows a builtin method to accept arguments that may be one of two
10731082
/// types, raising a `TypeError` if it is neither.
10741083
///

vm/src/stdlib/io.rs

Lines changed: 30 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,19 @@ use std::io::prelude::*;
66
use std::io::Cursor;
77
use std::io::SeekFrom;
88

9-
use num_bigint::ToBigInt;
109
use num_traits::ToPrimitive;
1110

1211
use super::os;
1312
use crate::function::{OptionalArg, PyFuncArgs};
1413
use crate::obj::objbytearray::PyByteArray;
1514
use crate::obj::objbytes;
1615
use crate::obj::objbytes::PyBytes;
17-
use crate::obj::objint;
18-
use crate::obj::objstr;
16+
use crate::obj::objint::{self, PyIntRef};
17+
use crate::obj::objstr::{self, PyStringRef};
1918
use crate::obj::objtype;
2019
use crate::obj::objtype::PyClassRef;
2120
use crate::pyobject::TypeProtocol;
22-
use crate::pyobject::{BufferProtocol, PyObjectRef, PyRef, PyResult, PyValue};
21+
use crate::pyobject::{BufferProtocol, Either, PyObjectRef, PyRef, PyResult, PyValue};
2322
use crate::vm::VirtualMachine;
2423

2524
fn byte_count(bytes: OptionalArg<Option<PyObjectRef>>) -> i64 {
@@ -102,7 +101,7 @@ impl PyValue for PyStringIO {
102101

103102
impl PyStringIORef {
104103
//write string to underlying vector
105-
fn write(self, data: objstr::PyStringRef, vm: &VirtualMachine) -> PyResult {
104+
fn write(self, data: PyStringRef, vm: &VirtualMachine) -> PyResult {
106105
let bytes = &data.value.clone().into_bytes();
107106

108107
match self.buffer.borrow_mut().write(bytes.to_vec()) {
@@ -312,32 +311,34 @@ fn compute_c_flag(mode: &str) -> u32 {
312311
flag as u32
313312
}
314313

315-
fn file_io_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
316-
arg_check!(
317-
vm,
318-
args,
319-
required = [(file_io, None), (name, None)],
320-
optional = [(mode, Some(vm.ctx.str_type()))]
321-
);
322-
323-
let file_no = if objtype::isinstance(&name, &vm.ctx.str_type()) {
324-
let rust_mode = mode.map_or("r".to_string(), objstr::get_value);
325-
let args = vec![
326-
name.clone(),
327-
vm.ctx
328-
.new_int(compute_c_flag(&rust_mode).to_bigint().unwrap()),
329-
];
330-
os::os_open(vm, PyFuncArgs::new(args, vec![]))?
331-
} else if objtype::isinstance(&name, &vm.ctx.int_type()) {
332-
name.clone()
333-
} else {
334-
return Err(vm.new_type_error("name parameter must be string or int".to_string()));
314+
fn file_io_init(
315+
file_io: PyObjectRef,
316+
name: Either<PyStringRef, PyIntRef>,
317+
mode: OptionalArg<PyStringRef>,
318+
vm: &VirtualMachine,
319+
) -> PyResult {
320+
let file_no = match &name {
321+
Either::A(name) => {
322+
let mode = match mode {
323+
OptionalArg::Present(mode) => compute_c_flag(mode.as_str()),
324+
OptionalArg::Missing => libc::O_RDONLY as _,
325+
};
326+
let fno = os::os_open(
327+
name.clone(),
328+
mode as _,
329+
OptionalArg::Missing,
330+
OptionalArg::Missing,
331+
vm,
332+
)?;
333+
vm.new_int(fno)
334+
}
335+
Either::B(fno) => fno.clone().into_object(),
335336
};
336337

337-
vm.set_attr(file_io, "name", name.clone())?;
338-
vm.set_attr(file_io, "fileno", file_no)?;
339-
vm.set_attr(file_io, "closefd", vm.new_bool(false))?;
340-
vm.set_attr(file_io, "closed", vm.new_bool(false))?;
338+
vm.set_attr(&file_io, "name", name.into_object())?;
339+
vm.set_attr(&file_io, "fileno", file_no)?;
340+
vm.set_attr(&file_io, "closefd", vm.new_bool(false))?;
341+
vm.set_attr(&file_io, "closed", vm.new_bool(false))?;
341342
Ok(vm.get_none())
342343
}
343344

vm/src/stdlib/os.rs

Lines changed: 48 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,13 @@ use nix::pty::openpty;
1919
use nix::unistd::{self, Gid, Pid, Uid, Whence};
2020
use num_traits::cast::ToPrimitive;
2121

22-
use crate::function::{IntoPyNativeFunc, PyFuncArgs};
22+
use crate::function::{IntoPyNativeFunc, OptionalArg, PyFuncArgs};
2323
use crate::obj::objbytes::PyBytesRef;
2424
use crate::obj::objdict::PyDictRef;
25-
use crate::obj::objint::{self, PyInt, PyIntRef};
25+
use crate::obj::objint::{self, PyIntRef};
2626
use crate::obj::objiter;
2727
use crate::obj::objset::PySet;
28-
use crate::obj::objstr::{self, PyString, PyStringRef};
28+
use crate::obj::objstr::{self, PyStringRef};
2929
use crate::obj::objtype::{self, PyClassRef};
3030
use crate::pyobject::{
3131
Either, ItemProtocol, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef,
@@ -95,59 +95,58 @@ pub fn os_close(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
9595
Ok(vm.get_none())
9696
}
9797

98-
#[cfg(any(unix, windows))]
99-
pub fn os_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
100-
arg_check!(
101-
vm,
102-
args,
103-
required = [
104-
(name, Some(vm.ctx.str_type())),
105-
(flags, Some(vm.ctx.int_type()))
106-
],
107-
optional = [
108-
(_mode, Some(vm.ctx.int_type())),
109-
(dir_fd, Some(vm.ctx.int_type()))
110-
]
111-
);
98+
#[cfg(unix)]
99+
type OpenFlags = i32;
100+
#[cfg(windows)]
101+
type OpenFlags = u32;
112102

113-
let name = name.clone().downcast::<PyString>().unwrap();
114-
let dir_fd = if let Some(obj) = dir_fd {
115-
DirFd {
116-
dir_fd: Some(obj.clone().downcast::<PyInt>().unwrap()),
117-
}
118-
} else {
119-
DirFd::default()
103+
#[cfg(any(unix, windows))]
104+
pub fn os_open(
105+
name: PyStringRef,
106+
flags: OpenFlags,
107+
_mode: OptionalArg<PyIntRef>,
108+
dir_fd: OptionalArg<PyIntRef>,
109+
vm: &VirtualMachine,
110+
) -> PyResult<i64> {
111+
let dir_fd = DirFd {
112+
dir_fd: dir_fd.into_option(),
120113
};
121-
let fname = &make_path(vm, name, &dir_fd).value;
114+
let fname = make_path(vm, name, &dir_fd);
122115

123-
let options = _set_file_model(&flags);
124-
let handle = options
125-
.open(&fname)
126-
.map_err(|err| convert_io_error(vm, err))?;
116+
let mut options = OpenOptions::new();
127117

128-
Ok(vm.ctx.new_int(raw_file_number(handle)))
129-
}
118+
macro_rules! bit_contains {
119+
($c:expr) => {
120+
flags & $c as OpenFlags == $c as OpenFlags
121+
};
122+
}
130123

131-
#[cfg(unix)]
132-
fn _set_file_model(flags: &PyObjectRef) -> OpenOptions {
133-
let flags = objint::get_value(flags).to_i32().unwrap();
134-
let mut options = OpenOptions::new();
135-
options.read(flags == libc::O_RDONLY);
136-
options.write(flags & libc::O_WRONLY != 0);
137-
options.append(flags & libc::O_APPEND != 0);
138-
options.custom_flags(flags);
139-
options
140-
}
124+
if bit_contains!(libc::O_RDWR) {
125+
options.read(true).write(true);
126+
} else if bit_contains!(libc::O_WRONLY) {
127+
options.write(true);
128+
} else if bit_contains!(libc::O_RDONLY) {
129+
options.read(true);
130+
}
131+
132+
if bit_contains!(libc::O_APPEND) {
133+
options.append(true);
134+
}
135+
136+
if bit_contains!(libc::O_CREAT) {
137+
if bit_contains!(libc::O_EXCL) {
138+
options.create_new(true);
139+
} else {
140+
options.create(true);
141+
}
142+
}
141143

142-
#[cfg(windows)]
143-
fn _set_file_model(flags: &PyObjectRef) -> OpenOptions {
144-
let flags = objint::get_value(flags).to_u32().unwrap();
145-
let mut options = OpenOptions::new();
146-
options.read((flags as i32) == libc::O_RDONLY);
147-
options.write((flags as i32) & libc::O_WRONLY != 0);
148-
options.append((flags as i32) & libc::O_APPEND != 0);
149144
options.custom_flags(flags);
150-
options
145+
let handle = options
146+
.open(fname.as_str())
147+
.map_err(|err| convert_io_error(vm, err))?;
148+
149+
Ok(raw_file_number(handle))
151150
}
152151

153152
#[cfg(all(not(unix), not(windows)))]

0 commit comments

Comments
 (0)