Skip to content

Commit dc7eb24

Browse files
committed
Support fileno as name in FileIO
1 parent 725cfcc commit dc7eb24

File tree

2 files changed

+34
-32
lines changed

2 files changed

+34
-32
lines changed

tests/snippets/stdlib_io.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from io import BufferedReader, FileIO
2+
import os
23

34
fi = FileIO('README.md')
45
bb = BufferedReader(fi)
@@ -14,3 +15,9 @@
1415
assert len(result) <= 8*1024
1516
assert len(result) >= 0
1617
assert isinstance(result, bytes)
18+
19+
fd = os.open('README.md', os.O_RDONLY)
20+
21+
with FileIO(fd) as fio:
22+
res2 = fio.read()
23+
assert res == res2

vm/src/stdlib/io.rs

Lines changed: 27 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@
22
* I/O core tools.
33
*/
44
use std::cell::RefCell;
5-
use std::fs::File;
65
use std::io::prelude::*;
7-
use std::io::BufReader;
86
use std::io::Cursor;
97
use std::io::SeekFrom;
108

@@ -20,6 +18,7 @@ use crate::obj::objbytearray::PyByteArray;
2018
use crate::obj::objbytes;
2119
use crate::obj::objint;
2220
use crate::obj::objstr;
21+
use crate::obj::objtype;
2322
use crate::obj::objtype::PyClassRef;
2423
use crate::pyobject::{BufferProtocol, PyObjectRef, PyRef, PyResult, PyValue};
2524
use crate::vm::VirtualMachine;
@@ -305,47 +304,43 @@ fn file_io_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
305304
arg_check!(
306305
vm,
307306
args,
308-
required = [(file_io, None), (name, Some(vm.ctx.str_type()))],
307+
required = [(file_io, None), (name, None)],
309308
optional = [(mode, Some(vm.ctx.str_type()))]
310309
);
311310

312-
let rust_mode = mode.map_or("r".to_string(), objstr::get_value);
313-
314-
match compute_c_flag(&rust_mode).to_bigint() {
315-
Some(os_mode) => {
316-
let args = vec![name.clone(), vm.ctx.new_int(os_mode)];
317-
let file_no = os::os_open(vm, PyFuncArgs::new(args, vec![]))?;
318-
319-
vm.set_attr(file_io, "name", name.clone())?;
320-
vm.set_attr(file_io, "fileno", file_no)?;
321-
vm.set_attr(file_io, "closefd", vm.new_bool(false))?;
322-
vm.set_attr(file_io, "closed", vm.new_bool(false))?;
311+
let file_no = if objtype::isinstance(&name, &vm.ctx.str_type()) {
312+
let rust_mode = mode.map_or("r".to_string(), objstr::get_value);
313+
let args = vec![
314+
name.clone(),
315+
vm.ctx
316+
.new_int(compute_c_flag(&rust_mode).to_bigint().unwrap()),
317+
];
318+
os::os_open(vm, PyFuncArgs::new(args, vec![]))?
319+
} else if objtype::isinstance(&name, &vm.ctx.int_type()) {
320+
name.clone()
321+
} else {
322+
return Err(vm.new_type_error("name parameter must be string or int".to_string()));
323+
};
323324

324-
Ok(vm.get_none())
325-
}
326-
None => Err(vm.new_type_error(format!("invalid mode {}", rust_mode))),
327-
}
325+
vm.set_attr(file_io, "name", name.clone())?;
326+
vm.set_attr(file_io, "fileno", file_no)?;
327+
vm.set_attr(file_io, "closefd", vm.new_bool(false))?;
328+
vm.set_attr(file_io, "closed", vm.new_bool(false))?;
329+
Ok(vm.get_none())
328330
}
329331

330332
fn file_io_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
331333
arg_check!(vm, args, required = [(file_io, None)]);
332-
let py_name = vm.get_attribute(file_io.clone(), "name")?;
333-
let f = match File::open(objstr::get_value(&py_name)) {
334-
Ok(v) => Ok(v),
335-
Err(_) => Err(vm.new_type_error("Error opening file".to_string())),
336-
};
337334

338-
let buffer = match f {
339-
Ok(v) => Ok(BufReader::new(v)),
340-
Err(_) => Err(vm.new_type_error("Error reading from file".to_string())),
341-
};
335+
let file_no = vm.get_attribute(file_io.clone(), "fileno")?;
336+
let raw_fd = objint::get_value(&file_no).to_i64().unwrap();
337+
338+
let mut handle = os::rust_file(raw_fd);
342339

343340
let mut bytes = vec![];
344-
if let Ok(mut buff) = buffer {
345-
match buff.read_to_end(&mut bytes) {
346-
Ok(_) => {}
347-
Err(_) => return Err(vm.new_value_error("Error reading from Buffer".to_string())),
348-
}
341+
match handle.read_to_end(&mut bytes) {
342+
Ok(_) => {}
343+
Err(_) => return Err(vm.new_value_error("Error reading from Buffer".to_string())),
349344
}
350345

351346
Ok(vm.ctx.new_bytes(bytes))

0 commit comments

Comments
 (0)