Merge pull request RustPython#1033 from palaviv/open-flags

Support more open flags

Author: coolreader18

Cargo.lock

@@ -4,12 +4,13 @@
 
 from testutils import assert_raises
 
-fd = os.open('README.md', 0)
+fd = os.open('README.md', os.O_RDONLY)
 assert fd > 0
 
 os.close(fd)
 assert_raises(OSError, lambda: os.read(fd, 10))
-assert_raises(FileNotFoundError, lambda: os.open('DOES_NOT_EXIST', 0))
+assert_raises(FileNotFoundError, lambda: os.open('DOES_NOT_EXIST', os.O_RDONLY))
+assert_raises(FileNotFoundError, lambda: os.open('DOES_NOT_EXIST', os.O_WRONLY))
 
 
 assert os.O_RDONLY == 0
@@ -88,14 +89,17 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 with TestWithTempDir() as tmpdir:
 	fname = os.path.join(tmpdir, FILE_NAME)
-	with open(fname, "wb"):
-		pass
-	fd = os.open(fname, 1)
+	fd = os.open(fname, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
 	assert os.write(fd, CONTENT2) == len(CONTENT2)
+	os.close(fd)
+
+	fd = os.open(fname, os.O_WRONLY | os.O_APPEND)
 	assert os.write(fd, CONTENT3) == len(CONTENT3)
 	os.close(fd)
 
-	fd = os.open(fname, 0)
+	assert_raises(FileExistsError, lambda: os.open(fname, os.O_WRONLY | os.O_CREAT | os.O_EXCL))
+
+	fd = os.open(fname, os.O_RDONLY)
 	assert os.read(fd, len(CONTENT2)) == CONTENT2
 	assert os.read(fd, len(CONTENT3)) == CONTENT3
 	os.close(fd)

tests/snippets/stdlib_os.py

@@ -32,6 +32,7 @@ indexmap = "1.0.2"
 crc = "^1.0.0"
 bincode = "1.1.4"
 unicode_categories = "0.1.1"
+bitflags = "1.1"
 
 
 # TODO: release and publish to crates.io

vm/Cargo.toml

@@ -856,6 +856,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
         "IndexError" => ctx.exceptions.index_error.clone(),
         "ImportError" => ctx.exceptions.import_error.clone(),
         "FileNotFoundError" => ctx.exceptions.file_not_found_error.clone(),
+        "FileExistsError" => ctx.exceptions.file_exists_error.clone(),
         "StopIteration" => ctx.exceptions.stop_iteration.clone(),
         "ZeroDivisionError" => ctx.exceptions.zero_division_error.clone(),
         "KeyError" => ctx.exceptions.key_error.clone(),

vm/src/builtins.rs

@@ -148,6 +148,7 @@ pub struct ExceptionZoo {
     pub base_exception_type: PyClassRef,
     pub exception_type: PyClassRef,
     pub file_not_found_error: PyClassRef,
+    pub file_exists_error: PyClassRef,
     pub import_error: PyClassRef,
     pub index_error: PyClassRef,
     pub key_error: PyClassRef,
@@ -203,6 +204,7 @@ impl ExceptionZoo {
         let not_implemented_error = create_type("NotImplementedError", &type_type, &runtime_error);
         let file_not_found_error = create_type("FileNotFoundError", &type_type, &os_error);
         let permission_error = create_type("PermissionError", &type_type, &os_error);
+        let file_exists_error = create_type("FileExistsError", &type_type, &os_error);
 
         let warning = create_type("Warning", &type_type, &exception_type);
         let bytes_warning = create_type("BytesWarning", &type_type, &warning);
@@ -224,6 +226,7 @@ impl ExceptionZoo {
             base_exception_type,
             exception_type,
             file_not_found_error,
+            file_exists_error,
             import_error,
             index_error,
             key_error,

vm/src/exceptions.rs

@@ -21,16 +21,6 @@ use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{BufferProtocol, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
-fn compute_c_flag(mode: &str) -> u16 {
-    match mode {
-        "w" => 512,
-        "x" => 512,
-        "a" => 8,
-        "+" => 2,
-        _ => 0,
-    }
-}
-
 #[derive(Debug)]
 struct PyStringIO {
     data: RefCell<String>,
@@ -132,6 +122,21 @@ fn buffered_reader_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.ctx.new_bytes(result))
 }
 
+fn compute_c_flag(mode: &str) -> u32 {
+    let flags = match mode {
+        "w" => os::FileCreationFlags::O_WRONLY | os::FileCreationFlags::O_CREAT,
+        "x" => {
+            os::FileCreationFlags::O_WRONLY
+                | os::FileCreationFlags::O_CREAT
+                | os::FileCreationFlags::O_EXCL
+        }
+        "a" => os::FileCreationFlags::O_APPEND,
+        "+" => os::FileCreationFlags::O_RDWR,
+        _ => os::FileCreationFlags::O_RDONLY,
+    };
+    flags.bits()
+}
+
 fn file_io_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,

vm/src/stdlib/io.rs

@@ -5,6 +5,7 @@ use std::io::{self, ErrorKind, Read, Write};
 use std::time::{Duration, SystemTime};
 use std::{env, fs};
 
+use bitflags::bitflags;
 use num_traits::cast::ToPrimitive;
 
 use crate::function::{IntoPyNativeFunc, PyFuncArgs};
@@ -81,15 +82,31 @@ pub fn os_close(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
 
+bitflags! {
+     pub struct FileCreationFlags: u32 {
+        // https://elixir.bootlin.com/linux/v4.8/source/include/uapi/asm-generic/fcntl.h
+        const O_RDONLY = 0o0000_0000;
+        const O_WRONLY = 0o0000_0001;
+        const O_RDWR = 0o0000_0002;
+        const O_CREAT = 0o0000_0100;
+        const O_EXCL = 0o0000_0200;
+        const O_APPEND = 0o0000_2000;
+        const O_NONBLOCK = 0o0000_4000;
+    }
+}
+
 pub fn os_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
         required = [
             (name, Some(vm.ctx.str_type())),
-            (mode, Some(vm.ctx.int_type()))
+            (flags, Some(vm.ctx.int_type()))
         ],
-        optional = [(dir_fd, Some(vm.ctx.int_type()))]
+        optional = [
+            (_mode, Some(vm.ctx.int_type())),
+            (dir_fd, Some(vm.ctx.int_type()))
+        ]
     );
 
     let name = name.clone().downcast::<PyString>().unwrap();
@@ -102,14 +119,32 @@ pub fn os_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     };
     let fname = &make_path(vm, name, &dir_fd).value;
 
-    let handle = match objint::get_value(mode).to_u16().unwrap() {
-        0 => OpenOptions::new().read(true).open(&fname),
-        1 => OpenOptions::new().write(true).open(&fname),
-        2 => OpenOptions::new().read(true).write(true).open(&fname),
-        512 => OpenOptions::new().write(true).create(true).open(&fname),
-        _ => OpenOptions::new().read(true).open(&fname),
+    let flags = FileCreationFlags::from_bits(objint::get_value(flags).to_u32().unwrap())
+        .ok_or(vm.new_value_error("Unsupported flag".to_string()))?;
+
+    let mut options = &mut OpenOptions::new();
+
+    if flags.contains(FileCreationFlags::O_WRONLY) {
+        options = options.write(true);
+    } else if flags.contains(FileCreationFlags::O_RDWR) {
+        options = options.read(true).write(true);
+    } else {
+        options = options.read(true);
     }
-    .map_err(|err| match err.kind() {
+
+    if flags.contains(FileCreationFlags::O_APPEND) {
+        options = options.append(true);
+    }
+
+    if flags.contains(FileCreationFlags::O_CREAT) {
+        if flags.contains(FileCreationFlags::O_EXCL) {
+            options = options.create_new(true);
+        } else {
+            options = options.create(true);
+        }
+    }
+
+    let handle = options.open(&fname).map_err(|err| match err.kind() {
         ErrorKind::NotFound => {
             let exc_type = vm.ctx.exceptions.file_not_found_error.clone();
             vm.new_exception(exc_type, format!("No such file or directory: {}", &fname))
@@ -118,6 +153,10 @@ pub fn os_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
             let exc_type = vm.ctx.exceptions.permission_error.clone();
             vm.new_exception(exc_type, format!("Permission denied: {}", &fname))
         }
+        ErrorKind::AlreadyExists => {
+            let exc_type = vm.ctx.exceptions.file_exists_error.clone();
+            vm.new_exception(exc_type, format!("File exists: {}", &fname))
+        }
         _ => vm.new_value_error("Unhandled file IO error".to_string()),
     })?;
 
@@ -743,12 +782,13 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "getcwd" => ctx.new_rustfunc(os_getcwd),
         "chdir" => ctx.new_rustfunc(os_chdir),
         "fspath" => ctx.new_rustfunc(os_fspath),
-        "O_RDONLY" => ctx.new_int(0),
-        "O_WRONLY" => ctx.new_int(1),
-        "O_RDWR" => ctx.new_int(2),
-        "O_NONBLOCK" => ctx.new_int(4),
-        "O_APPEND" => ctx.new_int(8),
-        "O_CREAT" => ctx.new_int(512)
+        "O_RDONLY" => ctx.new_int(FileCreationFlags::O_RDONLY.bits()),
+        "O_WRONLY" => ctx.new_int(FileCreationFlags::O_WRONLY.bits()),
+        "O_RDWR" => ctx.new_int(FileCreationFlags::O_RDWR.bits()),
+        "O_NONBLOCK" => ctx.new_int(FileCreationFlags::O_NONBLOCK.bits()),
+        "O_APPEND" => ctx.new_int(FileCreationFlags::O_APPEND.bits()),
+        "O_EXCL" => ctx.new_int(FileCreationFlags::O_EXCL.bits()),
+        "O_CREAT" => ctx.new_int(FileCreationFlags::O_CREAT.bits())
     });
 
     for support in support_funcs {