rustpython_vm/function/
fspath.rs

1use crate::{
2    builtins::{PyBytes, PyBytesRef, PyStrRef},
3    convert::{IntoPyException, ToPyObject},
4    function::PyStr,
5    protocol::PyBuffer,
6    PyObjectRef, PyResult, TryFromObject, VirtualMachine,
7};
8use std::{ffi::OsStr, path::PathBuf};
9
10#[derive(Clone)]
11pub enum FsPath {
12    Str(PyStrRef),
13    Bytes(PyBytesRef),
14}
15
16impl FsPath {
17    // PyOS_FSPath in CPython
18    pub fn try_from(obj: PyObjectRef, check_for_nul: bool, vm: &VirtualMachine) -> PyResult<Self> {
19        let check_nul = |b: &[u8]| {
20            if !check_for_nul || memchr::memchr(b'\0', b).is_none() {
21                Ok(())
22            } else {
23                Err(crate::exceptions::cstring_error(vm))
24            }
25        };
26        let match1 = |obj: PyObjectRef| {
27            let pathlike = match_class!(match obj {
28                s @ PyStr => {
29                    check_nul(s.as_str().as_bytes())?;
30                    FsPath::Str(s)
31                }
32                b @ PyBytes => {
33                    check_nul(&b)?;
34                    FsPath::Bytes(b)
35                }
36                obj => return Ok(Err(obj)),
37            });
38            Ok(Ok(pathlike))
39        };
40        let obj = match match1(obj)? {
41            Ok(pathlike) => return Ok(pathlike),
42            Err(obj) => obj,
43        };
44        let method =
45            vm.get_method_or_type_error(obj.clone(), identifier!(vm, __fspath__), || {
46                format!(
47                    "should be string, bytes, os.PathLike or integer, not {}",
48                    obj.class().name()
49                )
50            })?;
51        let result = method.call((), vm)?;
52        match1(result)?.map_err(|result| {
53            vm.new_type_error(format!(
54                "expected {}.__fspath__() to return str or bytes, not {}",
55                obj.class().name(),
56                result.class().name(),
57            ))
58        })
59    }
60
61    pub fn as_os_str(&self, vm: &VirtualMachine) -> PyResult<&OsStr> {
62        // TODO: FS encodings
63        match self {
64            FsPath::Str(s) => Ok(s.as_str().as_ref()),
65            FsPath::Bytes(b) => Self::bytes_as_osstr(b.as_bytes(), vm),
66        }
67    }
68
69    pub fn as_bytes(&self) -> &[u8] {
70        // TODO: FS encodings
71        match self {
72            FsPath::Str(s) => s.as_str().as_bytes(),
73            FsPath::Bytes(b) => b.as_bytes(),
74        }
75    }
76
77    pub fn as_str(&self) -> &str {
78        match self {
79            FsPath::Bytes(b) => std::str::from_utf8(b).unwrap(),
80            FsPath::Str(s) => s.as_str(),
81        }
82    }
83
84    pub fn to_path_buf(&self, vm: &VirtualMachine) -> PyResult<PathBuf> {
85        let path = match self {
86            FsPath::Str(s) => PathBuf::from(s.as_str()),
87            FsPath::Bytes(b) => PathBuf::from(Self::bytes_as_osstr(b, vm)?),
88        };
89        Ok(path)
90    }
91
92    pub fn to_cstring(&self, vm: &VirtualMachine) -> PyResult<std::ffi::CString> {
93        std::ffi::CString::new(self.as_bytes()).map_err(|e| e.into_pyexception(vm))
94    }
95
96    #[cfg(windows)]
97    pub fn to_widecstring(&self, vm: &VirtualMachine) -> PyResult<widestring::WideCString> {
98        widestring::WideCString::from_os_str(self.as_os_str(vm)?)
99            .map_err(|err| err.into_pyexception(vm))
100    }
101
102    pub fn bytes_as_osstr<'a>(b: &'a [u8], vm: &VirtualMachine) -> PyResult<&'a std::ffi::OsStr> {
103        rustpython_common::os::bytes_as_osstr(b)
104            .map_err(|_| vm.new_unicode_decode_error("can't decode path for utf-8".to_owned()))
105    }
106}
107
108impl ToPyObject for FsPath {
109    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
110        match self {
111            Self::Str(s) => s.into(),
112            Self::Bytes(b) => b.into(),
113        }
114    }
115}
116
117impl TryFromObject for FsPath {
118    // PyUnicode_FSDecoder in CPython
119    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
120        let obj = match obj.try_to_value::<PyBuffer>(vm) {
121            Ok(buffer) => {
122                let mut bytes = vec![];
123                buffer.append_to(&mut bytes);
124                vm.ctx.new_bytes(bytes).into()
125            }
126            Err(_) => obj,
127        };
128        Self::try_from(obj, true, vm)
129    }
130}
Morty Proxy This is a proxified and sanitized view of the page, visit original site.