Skip to content

Commit 55e273f

Browse files
committed
Overhaul PyPathLike/OutputMode a *tiny* bit
1 parent 642da4c commit 55e273f

File tree

2 files changed

+95
-86
lines changed

2 files changed

+95
-86
lines changed

vm/src/stdlib/os.rs

Lines changed: 94 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use std::io::{self, ErrorKind, Read, Write};
77
use std::os::unix::fs::OpenOptionsExt;
88
#[cfg(windows)]
99
use std::os::windows::fs::OpenOptionsExt;
10+
use std::path::PathBuf;
1011
use std::time::{Duration, SystemTime};
1112
use std::{env, fs};
1213

@@ -119,75 +120,95 @@ enum OutputMode {
119120
Bytes,
120121
}
121122

122-
fn output_by_mode(val: String, mode: OutputMode, vm: &VirtualMachine) -> PyObjectRef {
123-
match mode {
124-
OutputMode::String => vm.ctx.new_str(val),
125-
OutputMode::Bytes => vm.ctx.new_bytes(val.into_bytes()),
123+
impl OutputMode {
124+
fn process_path(self, path: impl Into<PathBuf>, vm: &VirtualMachine) -> PyResult {
125+
fn inner(mode: OutputMode, path: PathBuf, vm: &VirtualMachine) -> PyResult {
126+
let path_as_string = |p: PathBuf| {
127+
p.into_os_string().into_string().map_err(|_| {
128+
vm.new_unicode_decode_error(
129+
"Can't convert OS path to valid UTF-8 string".into(),
130+
)
131+
})
132+
};
133+
match mode {
134+
OutputMode::String => path_as_string(path).map(|s| vm.new_str(s)),
135+
OutputMode::Bytes => {
136+
#[cfg(unix)]
137+
{
138+
use std::os::unix::ffi::OsStringExt;
139+
Ok(vm.ctx.new_bytes(path.into_os_string().into_vec()))
140+
}
141+
#[cfg(windows)]
142+
{
143+
path_as_string(path).map(|s| vm.ctx.new_bytes(s.into_bytes()))
144+
}
145+
}
146+
}
147+
}
148+
inner(self, path.into(), vm)
126149
}
127150
}
128151

129152
pub struct PyPathLike {
130-
pub path: ffi::OsString,
153+
pub path: PathBuf,
131154
mode: OutputMode,
132155
}
133156

134157
impl PyPathLike {
135158
pub fn new_str(path: String) -> Self {
136159
PyPathLike {
137-
path: ffi::OsString::from(path),
160+
path: PathBuf::from(path),
138161
mode: OutputMode::String,
139162
}
140163
}
141164
#[cfg(windows)]
142165
pub fn wide(&self) -> Vec<u16> {
143166
use std::os::windows::ffi::OsStrExt;
144-
self.path.encode_wide().chain(std::iter::once(0)).collect()
167+
self.path
168+
.as_os_str()
169+
.encode_wide()
170+
.chain(std::iter::once(0))
171+
.collect()
172+
}
173+
#[cfg(unix)]
174+
pub fn into_bytes(self) -> Vec<u8> {
175+
use std::os::unix::ffi::OsStringExt;
176+
self.path.into_os_string().into_vec()
145177
}
146178
}
147179

148180
impl TryFromObject for PyPathLike {
149181
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
150-
match_class!(match obj.clone() {
151-
l @ PyString => {
152-
Ok(PyPathLike {
153-
path: ffi::OsString::from(l.as_str()),
182+
let match1 = |obj: &PyObjectRef| {
183+
let pathlike = match_class!(match obj {
184+
ref l @ PyString => PyPathLike {
185+
path: l.as_str().into(),
154186
mode: OutputMode::String,
155-
})
156-
}
157-
i @ PyBytes => {
158-
Ok(PyPathLike {
159-
path: bytes_as_osstr(&i, vm)?.to_os_string(),
187+
},
188+
ref i @ PyBytes => PyPathLike {
189+
path: bytes_as_osstr(&i, vm)?.to_os_string().into(),
160190
mode: OutputMode::Bytes,
161-
})
162-
}
163-
obj => {
164-
let method = vm.get_method_or_type_error(obj.clone(), "__fspath__", || {
165-
format!(
166-
"expected str, bytes or os.PathLike object, not '{}'",
167-
obj.class().name
168-
)
169-
})?;
170-
let result = vm.invoke(&method, PyFuncArgs::default())?;
171-
match_class!(match result.clone() {
172-
l @ PyString => {
173-
Ok(PyPathLike {
174-
path: ffi::OsString::from(l.as_str()),
175-
mode: OutputMode::String,
176-
})
177-
}
178-
i @ PyBytes => {
179-
Ok(PyPathLike {
180-
path: bytes_as_osstr(&i, vm)?.to_os_string(),
181-
mode: OutputMode::Bytes,
182-
})
183-
}
184-
_ => Err(vm.new_type_error(format!(
185-
"expected {}.__fspath__() to return str or bytes, not '{}'",
186-
obj.class().name,
187-
result.class().name
188-
))),
189-
})
190-
}
191+
},
192+
_ => return Ok(None),
193+
});
194+
Ok(Some(pathlike))
195+
};
196+
if let Some(pathlike) = match1(&obj)? {
197+
return Ok(pathlike);
198+
}
199+
let method = vm.get_method_or_type_error(obj.clone(), "__fspath__", || {
200+
format!(
201+
"expected str, bytes or os.PathLike object, not '{}'",
202+
obj.class().name
203+
)
204+
})?;
205+
let result = vm.invoke(&method, PyFuncArgs::default())?;
206+
match1(&result)?.ok_or_else(|| {
207+
vm.new_type_error(format!(
208+
"expected {}.__fspath__() to return str or bytes, not '{}'",
209+
obj.class().name,
210+
result.class().name,
211+
))
191212
})
192213
}
193214
}
@@ -529,31 +550,26 @@ fn os_listdir(path: PyPathLike, vm: &VirtualMachine) -> PyResult {
529550
let res: PyResult<Vec<PyObjectRef>> = fs::read_dir(&path.path)
530551
.map_err(|err| convert_io_error(vm, err))?
531552
.map(|entry| match entry {
532-
Ok(entry_path) => Ok(output_by_mode(
533-
entry_path.file_name().into_string().unwrap(),
534-
path.mode,
535-
vm,
536-
)),
553+
Ok(entry_path) => path.mode.process_path(entry_path.file_name(), vm),
537554
Err(s) => Err(convert_io_error(vm, s)),
538555
})
539556
.collect();
540557
Ok(vm.ctx.new_list(res?))
541558
}
542559

543560
fn bytes_as_osstr<'a>(b: &'a [u8], vm: &VirtualMachine) -> PyResult<&'a ffi::OsStr> {
544-
let os_str = {
545-
#[cfg(unix)]
546-
{
547-
use std::os::unix::ffi::OsStrExt;
548-
Some(ffi::OsStr::from_bytes(b))
549-
}
550-
#[cfg(not(unix))]
551-
{
552-
std::str::from_utf8(b).ok().map(|s| s.as_ref())
553-
}
554-
};
555-
os_str
556-
.ok_or_else(|| vm.new_value_error("Can't convert bytes to str for env function".to_owned()))
561+
#[cfg(unix)]
562+
{
563+
let _ = vm;
564+
use std::os::unix::ffi::OsStrExt;
565+
Ok(ffi::OsStr::from_bytes(b))
566+
}
567+
#[cfg(not(unix))]
568+
{
569+
std::str::from_utf8(b).map(|s| s.as_ref()).map_err(|_| {
570+
vm.new_value_error("Can't convert bytes to str for env function".to_owned())
571+
})
572+
}
557573
}
558574

559575
fn os_putenv(
@@ -609,12 +625,10 @@ fn _os_environ(vm: &VirtualMachine) -> PyDictRef {
609625
}
610626

611627
fn os_readlink(path: PyPathLike, dir_fd: DirFd, vm: &VirtualMachine) -> PyResult {
628+
let mode = path.mode;
612629
let path = make_path(vm, &path, &dir_fd);
613630
let path = fs::read_link(path).map_err(|err| convert_io_error(vm, err))?;
614-
let path = path.into_os_string().into_string().map_err(|_osstr| {
615-
vm.new_unicode_decode_error("Can't convert OS path to valid UTF-8 string".into())
616-
})?;
617-
Ok(vm.ctx.new_str(path))
631+
mode.process_path(path, vm)
618632
}
619633

620634
#[derive(Debug)]
@@ -644,14 +658,12 @@ struct FollowSymlinks {
644658
}
645659

646660
impl DirEntryRef {
647-
fn name(self, vm: &VirtualMachine) -> PyObjectRef {
648-
let file_name = self.entry.file_name().into_string().unwrap();
649-
output_by_mode(file_name, self.mode, vm)
661+
fn name(self, vm: &VirtualMachine) -> PyResult {
662+
self.mode.process_path(self.entry.file_name(), vm)
650663
}
651664

652-
fn path(self, vm: &VirtualMachine) -> PyObjectRef {
653-
let path = self.entry.path().to_str().unwrap().to_owned();
654-
output_by_mode(path, self.mode, vm)
665+
fn path(self, vm: &VirtualMachine) -> PyResult {
666+
self.mode.process_path(self.entry.path(), vm)
655667
}
656668

657669
#[allow(clippy::match_bool)]
@@ -696,7 +708,7 @@ impl DirEntryRef {
696708
fn stat(self, dir_fd: DirFd, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult {
697709
os_stat(
698710
Either::A(PyPathLike {
699-
path: self.entry.path().into_os_string(),
711+
path: self.entry.path(),
700712
mode: OutputMode::String,
701713
}),
702714
dir_fd,
@@ -1174,8 +1186,8 @@ fn os_chmod(
11741186
Ok(())
11751187
}
11761188

1177-
fn os_fspath(path: PyPathLike, vm: &VirtualMachine) -> PyObjectRef {
1178-
output_by_mode(path.path.to_str().unwrap().to_owned(), path.mode, vm)
1189+
fn os_fspath(path: PyPathLike, vm: &VirtualMachine) -> PyResult {
1190+
path.mode.process_path(path.path, vm)
11791191
}
11801192

11811193
fn os_rename(src: PyPathLike, dst: PyPathLike, vm: &VirtualMachine) -> PyResult<()> {
@@ -1643,11 +1655,10 @@ fn os_setgroups(group_ids: PyIterable<u32>, vm: &VirtualMachine) -> PyResult<()>
16431655

16441656
#[cfg(unix)]
16451657
fn envp_from_dict(dict: PyDictRef, vm: &VirtualMachine) -> PyResult<Vec<ffi::CString>> {
1646-
use std::os::unix::ffi::OsStringExt;
16471658
dict.into_iter()
16481659
.map(|(k, v)| {
1649-
let k = PyPathLike::try_from_object(vm, k)?.path.into_vec();
1650-
let v = PyPathLike::try_from_object(vm, v)?.path.into_vec();
1660+
let k = PyPathLike::try_from_object(vm, k)?.into_bytes();
1661+
let v = PyPathLike::try_from_object(vm, v)?.into_bytes();
16511662
if k.contains(&0) {
16521663
return Err(
16531664
vm.new_value_error("envp dict key cannot contain a nul byte".to_owned())
@@ -1698,8 +1709,7 @@ enum PosixSpawnFileActionIdentifier {
16981709
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
16991710
impl PosixSpawnArgs {
17001711
fn spawn(self, spawnp: bool, vm: &VirtualMachine) -> PyResult<libc::pid_t> {
1701-
use std::os::unix::ffi::OsStringExt;
1702-
let path = ffi::CString::new(self.path.path.into_vec())
1712+
let path = ffi::CString::new(self.path.into_bytes())
17031713
.map_err(|_| vm.new_value_error("path should not have nul bytes".to_owned()))?;
17041714

17051715
let mut file_actions = unsafe {
@@ -1722,7 +1732,7 @@ impl PosixSpawnArgs {
17221732
let ret = match id {
17231733
PosixSpawnFileActionIdentifier::Open => {
17241734
let (fd, path, oflag, mode): (_, PyPathLike, _, _) = args.bind(vm)?;
1725-
let path = ffi::CString::new(path.path.into_vec()).map_err(|_| {
1735+
let path = ffi::CString::new(path.into_bytes()).map_err(|_| {
17261736
vm.new_value_error(
17271737
"POSIX_SPAWN_OPEN path should not have nul bytes".to_owned(),
17281738
)
@@ -1776,7 +1786,7 @@ impl PosixSpawnArgs {
17761786
.args
17771787
.iter(vm)?
17781788
.map(|res| {
1779-
ffi::CString::new(res?.path.into_vec())
1789+
ffi::CString::new(res?.into_bytes())
17801790
.map_err(|_| vm.new_value_error("path should not have nul bytes".to_owned()))
17811791
})
17821792
.collect::<Result<_, _>>()?;

vm/src/stdlib/posixsubprocess.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ use nix::{dir, errno::Errno, fcntl, unistd};
22
use std::convert::Infallible as Never;
33
use std::ffi::{CStr, CString};
44
use std::io::{self, prelude::*};
5-
use std::os::unix::ffi::OsStringExt;
65

76
use crate::pyobject::{PyObjectRef, PyResult, PySequence, TryFromObject};
87
use crate::VirtualMachine;
@@ -23,7 +22,7 @@ struct CStrPathLike {
2322
}
2423
impl TryFromObject for CStrPathLike {
2524
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
26-
let s = PyPathLike::try_from_object(vm, obj)?.path.into_vec();
25+
let s = PyPathLike::try_from_object(vm, obj)?.into_bytes();
2726
let s = CString::new(s)
2827
.map_err(|_| vm.new_value_error("embedded null character".to_owned()))?;
2928
Ok(CStrPathLike { s })

0 commit comments

Comments
 (0)