Skip to main content

rustpython_vm/function/
fspath.rs

1use crate::{
2    PyObjectRef, PyResult, TryFromObject, VirtualMachine,
3    builtins::{PyBytes, PyBytesRef, PyStrRef},
4    convert::{IntoPyException, ToPyObject},
5    function::PyStr,
6    protocol::PyBuffer,
7};
8use alloc::borrow::Cow;
9use std::{ffi::OsStr, path::PathBuf};
10
11/// Helper to implement os.fspath()
12#[derive(Clone)]
13pub enum FsPath {
14    Str(PyStrRef),
15    Bytes(PyBytesRef),
16}
17
18impl FsPath {
19    pub fn try_from_path_like(
20        obj: PyObjectRef,
21        check_for_nul: bool,
22        vm: &VirtualMachine,
23    ) -> PyResult<Self> {
24        Self::try_from(
25            obj,
26            check_for_nul,
27            "expected str, bytes or os.PathLike object",
28            vm,
29        )
30    }
31
32    // PyOS_FSPath
33    pub fn try_from(
34        obj: PyObjectRef,
35        check_for_nul: bool,
36        msg: &'static str,
37        vm: &VirtualMachine,
38    ) -> PyResult<Self> {
39        let check_nul = |b: &[u8]| {
40            if !check_for_nul || memchr::memchr(b'\0', b).is_none() {
41                Ok(())
42            } else {
43                Err(crate::exceptions::cstring_error(vm))
44            }
45        };
46        let match1 = |obj: PyObjectRef| {
47            let pathlike = match_class!(match obj {
48                s @ PyStr => {
49                    check_nul(s.as_bytes())?;
50                    Self::Str(s)
51                }
52                b @ PyBytes => {
53                    check_nul(&b)?;
54                    Self::Bytes(b)
55                }
56                obj => return Ok(Err(obj)),
57            });
58            Ok(Ok(pathlike))
59        };
60        let obj = match match1(obj)? {
61            Ok(pathlike) => return Ok(pathlike),
62            Err(obj) => obj,
63        };
64        let not_pathlike_error = || format!("{msg}, not {}", obj.class().name());
65        let method = vm.get_method_or_type_error(
66            obj.clone(),
67            identifier!(vm, __fspath__),
68            not_pathlike_error,
69        )?;
70        // If __fspath__ is explicitly set to None, treat it as if it doesn't have __fspath__
71        if vm.is_none(&method) {
72            return Err(vm.new_type_error(not_pathlike_error()));
73        }
74        let result = method.call((), vm)?;
75        match1(result)?.map_err(|result| {
76            vm.new_type_error(format!(
77                "expected {}.__fspath__() to return str or bytes, not {}",
78                obj.class().name(),
79                result.class().name(),
80            ))
81        })
82    }
83
84    pub fn as_os_str(&self, vm: &VirtualMachine) -> PyResult<Cow<'_, OsStr>> {
85        // TODO: FS encodings
86        match self {
87            Self::Str(s) => vm.fsencode(s),
88            Self::Bytes(b) => Self::bytes_as_os_str(b.as_bytes(), vm).map(Cow::Borrowed),
89        }
90    }
91
92    pub fn as_bytes(&self) -> &[u8] {
93        // TODO: FS encodings
94        match self {
95            Self::Str(s) => s.as_bytes(),
96            Self::Bytes(b) => b.as_bytes(),
97        }
98    }
99
100    pub fn to_string_lossy(&self) -> Cow<'_, str> {
101        match self {
102            Self::Str(s) => s.to_string_lossy(),
103            Self::Bytes(s) => String::from_utf8_lossy(s),
104        }
105    }
106
107    pub fn to_path_buf(&self, vm: &VirtualMachine) -> PyResult<PathBuf> {
108        let path = match self {
109            Self::Str(s) => PathBuf::from(vm.fsencode(s)?.as_ref() as &OsStr),
110            Self::Bytes(b) => PathBuf::from(Self::bytes_as_os_str(b, vm)?),
111        };
112        Ok(path)
113    }
114
115    pub fn to_cstring(&self, vm: &VirtualMachine) -> PyResult<alloc::ffi::CString> {
116        alloc::ffi::CString::new(self.as_bytes()).map_err(|e| e.into_pyexception(vm))
117    }
118
119    #[cfg(windows)]
120    pub fn to_wide_cstring(&self, vm: &VirtualMachine) -> PyResult<widestring::WideCString> {
121        widestring::WideCString::from_os_str(self.as_os_str(vm)?)
122            .map_err(|err| err.into_pyexception(vm))
123    }
124
125    pub fn bytes_as_os_str<'a>(b: &'a [u8], vm: &VirtualMachine) -> PyResult<&'a std::ffi::OsStr> {
126        rustpython_common::os::bytes_as_os_str(b)
127            .map_err(|_| vm.new_unicode_decode_error("can't decode path for utf-8"))
128    }
129}
130
131impl ToPyObject for FsPath {
132    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
133        match self {
134            Self::Str(s) => s.into(),
135            Self::Bytes(b) => b.into(),
136        }
137    }
138}
139
140impl TryFromObject for FsPath {
141    // PyUnicode_FSDecoder in CPython
142    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
143        let obj = match obj.try_to_value::<PyBuffer>(vm) {
144            Ok(buffer) => {
145                let mut bytes = vec![];
146                buffer.append_to(&mut bytes);
147                vm.ctx.new_bytes(bytes).into()
148            }
149            Err(_) => obj,
150        };
151        Self::try_from_path_like(obj, true, vm)
152    }
153}