1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use crate::{
    builtins::{PyBytes, PyBytesRef, PyStrRef},
    convert::{IntoPyException, ToPyObject},
    function::PyStr,
    protocol::PyBuffer,
    PyObjectRef, PyResult, TryFromObject, VirtualMachine,
};
use std::{ffi::OsStr, path::PathBuf};

#[derive(Clone)]
pub enum FsPath {
    Str(PyStrRef),
    Bytes(PyBytesRef),
}

impl FsPath {
    // PyOS_FSPath in CPython
    pub fn try_from(obj: PyObjectRef, check_for_nul: bool, vm: &VirtualMachine) -> PyResult<Self> {
        let check_nul = |b: &[u8]| {
            if !check_for_nul || memchr::memchr(b'\0', b).is_none() {
                Ok(())
            } else {
                Err(crate::exceptions::cstring_error(vm))
            }
        };
        let match1 = |obj: PyObjectRef| {
            let pathlike = match_class!(match obj {
                s @ PyStr => {
                    check_nul(s.as_str().as_bytes())?;
                    FsPath::Str(s)
                }
                b @ PyBytes => {
                    check_nul(&b)?;
                    FsPath::Bytes(b)
                }
                obj => return Ok(Err(obj)),
            });
            Ok(Ok(pathlike))
        };
        let obj = match match1(obj)? {
            Ok(pathlike) => return Ok(pathlike),
            Err(obj) => obj,
        };
        let method =
            vm.get_method_or_type_error(obj.clone(), identifier!(vm, __fspath__), || {
                format!(
                    "should be string, bytes, os.PathLike or integer, not {}",
                    obj.class().name()
                )
            })?;
        let result = method.call((), vm)?;
        match1(result)?.map_err(|result| {
            vm.new_type_error(format!(
                "expected {}.__fspath__() to return str or bytes, not {}",
                obj.class().name(),
                result.class().name(),
            ))
        })
    }

    pub fn as_os_str(&self, vm: &VirtualMachine) -> PyResult<&OsStr> {
        // TODO: FS encodings
        match self {
            FsPath::Str(s) => Ok(s.as_str().as_ref()),
            FsPath::Bytes(b) => Self::bytes_as_osstr(b.as_bytes(), vm),
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        // TODO: FS encodings
        match self {
            FsPath::Str(s) => s.as_str().as_bytes(),
            FsPath::Bytes(b) => b.as_bytes(),
        }
    }

    pub fn as_str(&self) -> &str {
        match self {
            FsPath::Bytes(b) => std::str::from_utf8(b).unwrap(),
            FsPath::Str(s) => s.as_str(),
        }
    }

    pub fn to_path_buf(&self, vm: &VirtualMachine) -> PyResult<PathBuf> {
        let path = match self {
            FsPath::Str(s) => PathBuf::from(s.as_str()),
            FsPath::Bytes(b) => PathBuf::from(Self::bytes_as_osstr(b, vm)?),
        };
        Ok(path)
    }

    pub fn to_cstring(&self, vm: &VirtualMachine) -> PyResult<std::ffi::CString> {
        std::ffi::CString::new(self.as_bytes()).map_err(|e| e.into_pyexception(vm))
    }

    #[cfg(windows)]
    pub fn to_widecstring(&self, vm: &VirtualMachine) -> PyResult<widestring::WideCString> {
        widestring::WideCString::from_os_str(self.as_os_str(vm)?)
            .map_err(|err| err.into_pyexception(vm))
    }

    pub fn bytes_as_osstr<'a>(b: &'a [u8], vm: &VirtualMachine) -> PyResult<&'a std::ffi::OsStr> {
        rustpython_common::os::bytes_as_osstr(b)
            .map_err(|_| vm.new_unicode_decode_error("can't decode path for utf-8".to_owned()))
    }
}

impl ToPyObject for FsPath {
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        match self {
            Self::Str(s) => s.into(),
            Self::Bytes(b) => b.into(),
        }
    }
}

impl TryFromObject for FsPath {
    // PyUnicode_FSDecoder in CPython
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let obj = match obj.try_to_value::<PyBuffer>(vm) {
            Ok(buffer) => {
                let mut bytes = vec![];
                buffer.append_to(&mut bytes);
                vm.ctx.new_bytes(bytes).into()
            }
            Err(_) => obj,
        };
        Self::try_from(obj, true, vm)
    }
}