oxidized_importer/
python_resource_collector.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5/*! Python functionality for resource collection. */
6
7use {
8    crate::{
9        conversion::{path_to_pathlib_path, pyobject_to_pathbuf},
10        python_resource_types::{
11            PythonExtensionModule, PythonModuleBytecode, PythonModuleSource,
12            PythonPackageDistributionResource, PythonPackageResource,
13        },
14        python_resources::resource_to_pyobject,
15    },
16    anyhow::Context,
17    pyo3::{
18        exceptions::{PyTypeError, PyValueError},
19        ffi as pyffi,
20        prelude::*,
21        types::{PyBytes, PyList, PyTuple},
22        AsPyPointer,
23    },
24    python_packaging::{
25        bytecode::BytecodeCompiler,
26        location::{AbstractResourceLocation, ConcreteResourceLocation},
27        resource_collection::{CompiledResourcesCollection, PythonResourceCollector},
28    },
29    std::{
30        cell::RefCell,
31        path::{Path, PathBuf},
32    },
33};
34
35#[pyclass(module = "oxidized_importer")]
36pub struct PyTempDir {
37    cleanup: Py<PyAny>,
38    path: PathBuf,
39}
40
41impl PyTempDir {
42    pub fn new(py: Python) -> PyResult<Self> {
43        let temp_dir = py
44            .import("tempfile")?
45            .getattr("TemporaryDirectory")?
46            .call0()?;
47        let cleanup = temp_dir.getattr("cleanup")?.into_py(py);
48        let path = pyobject_to_pathbuf(py, temp_dir.getattr("name")?)?;
49
50        Ok(Self { cleanup, path })
51    }
52
53    pub fn path(&self) -> &Path {
54        &self.path
55    }
56}
57
58impl Drop for PyTempDir {
59    fn drop(&mut self) {
60        Python::with_gil(|py| {
61            if self.cleanup.call0(py).is_err() {
62                let cleanup = self.cleanup.as_ptr();
63                unsafe { pyffi::PyErr_WriteUnraisable(cleanup) }
64            }
65        });
66    }
67}
68
69#[pyclass(module = "oxidized_importer")]
70pub(crate) struct OxidizedResourceCollector {
71    collector: RefCell<PythonResourceCollector>,
72}
73
74#[pymethods]
75impl OxidizedResourceCollector {
76    fn __repr__(&self) -> &'static str {
77        "<OxidizedResourceCollector>"
78    }
79
80    #[new]
81    fn new(allowed_locations: Vec<String>) -> PyResult<Self> {
82        let allowed_locations = allowed_locations
83            .iter()
84            .map(|location| AbstractResourceLocation::try_from(location.as_str()))
85            .collect::<Result<Vec<_>, _>>()
86            .map_err(PyValueError::new_err)?;
87
88        let collector =
89            PythonResourceCollector::new(allowed_locations.clone(), allowed_locations, true, true);
90
91        Ok(Self {
92            collector: RefCell::new(collector),
93        })
94    }
95
96    #[getter]
97    fn allowed_locations<'p>(&self, py: Python<'p>) -> PyResult<&'p PyList> {
98        let values = self
99            .collector
100            .borrow()
101            .allowed_locations()
102            .iter()
103            .map(|l| l.to_string().into_py(py))
104            .collect::<Vec<Py<PyAny>>>();
105
106        Ok(PyList::new(py, &values))
107    }
108
109    fn add_in_memory(&self, resource: &PyAny) -> PyResult<()> {
110        let mut collector = self.collector.borrow_mut();
111        let typ = resource.get_type();
112        let repr = resource.repr()?;
113
114        match typ.name()? {
115            "PythonExtensionModule" => {
116                let module_cell = resource.cast_as::<PyCell<PythonExtensionModule>>()?;
117                let module = module_cell.borrow();
118                let resource = module.get_resource();
119
120                if resource.shared_library.is_some() {
121                    collector
122                        .add_python_extension_module(&resource, &ConcreteResourceLocation::InMemory)
123                        .with_context(|| format!("adding {}", repr))
124                        .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
125
126                    Ok(())
127                } else {
128                    Err(PyValueError::new_err(
129                        "PythonExtensionModule lacks a shared library",
130                    ))
131                }
132            }
133            "PythonModuleBytecode" => {
134                let module = resource.cast_as::<PyCell<PythonModuleBytecode>>()?;
135                collector
136                    .add_python_module_bytecode(
137                        &module.borrow().get_resource(),
138                        &ConcreteResourceLocation::InMemory,
139                    )
140                    .with_context(|| format!("adding {}", repr))
141                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
142
143                Ok(())
144            }
145            "PythonModuleSource" => {
146                let module = resource.cast_as::<PyCell<PythonModuleSource>>()?;
147                collector
148                    .add_python_module_source(
149                        &module.borrow().get_resource(),
150                        &ConcreteResourceLocation::InMemory,
151                    )
152                    .with_context(|| format!("adding {}", repr))
153                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
154
155                Ok(())
156            }
157            "PythonPackageResource" => {
158                let resource = resource.cast_as::<PyCell<PythonPackageResource>>()?;
159                collector
160                    .add_python_package_resource(
161                        &resource.borrow().get_resource(),
162                        &ConcreteResourceLocation::InMemory,
163                    )
164                    .with_context(|| format!("adding {}", repr))
165                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
166
167                Ok(())
168            }
169            "PythonPackageDistributionResource" => {
170                let resource = resource.cast_as::<PyCell<PythonPackageDistributionResource>>()?;
171                collector
172                    .add_python_package_distribution_resource(
173                        &resource.borrow().get_resource(),
174                        &ConcreteResourceLocation::InMemory,
175                    )
176                    .with_context(|| format!("adding {}", repr))
177                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
178
179                Ok(())
180            }
181            type_name => Err(PyTypeError::new_err(format!(
182                "cannot operate on {} values",
183                type_name
184            ))),
185        }
186    }
187
188    fn add_filesystem_relative(&self, prefix: String, resource: &PyAny) -> PyResult<()> {
189        let mut collector = self.collector.borrow_mut();
190
191        let repr = resource.repr()?;
192
193        match resource.get_type().name()? {
194            "PythonExtensionModule" => {
195                let module_cell = resource.cast_as::<PyCell<PythonExtensionModule>>()?;
196                let module = module_cell.borrow();
197                let resource = module.get_resource();
198
199                collector
200                    .add_python_extension_module(
201                        &resource,
202                        &ConcreteResourceLocation::RelativePath(prefix),
203                    )
204                    .with_context(|| format!("adding {}", repr))
205                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
206
207                Ok(())
208            }
209            "PythonModuleBytecode" => {
210                let module = resource.cast_as::<PyCell<PythonModuleBytecode>>()?;
211
212                collector
213                    .add_python_module_bytecode(
214                        &module.borrow().get_resource(),
215                        &ConcreteResourceLocation::RelativePath(prefix),
216                    )
217                    .with_context(|| format!("adding {}", repr))
218                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
219
220                Ok(())
221            }
222            "PythonModuleSource" => {
223                let module = resource.cast_as::<PyCell<PythonModuleSource>>()?;
224
225                collector
226                    .add_python_module_source(
227                        &module.borrow().get_resource(),
228                        &ConcreteResourceLocation::RelativePath(prefix),
229                    )
230                    .with_context(|| format!("adding {}", repr))
231                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
232
233                Ok(())
234            }
235            "PythonPackageResource" => {
236                let resource = resource.cast_as::<PyCell<PythonPackageResource>>()?;
237
238                collector
239                    .add_python_package_resource(
240                        &resource.borrow().get_resource(),
241                        &ConcreteResourceLocation::RelativePath(prefix),
242                    )
243                    .with_context(|| format!("adding {}", repr))
244                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
245
246                Ok(())
247            }
248            "PythonPackageDistributionResource" => {
249                let resource = resource.cast_as::<PyCell<PythonPackageDistributionResource>>()?;
250
251                collector
252                    .add_python_package_distribution_resource(
253                        &resource.borrow().get_resource(),
254                        &ConcreteResourceLocation::RelativePath(prefix),
255                    )
256                    .with_context(|| format!("adding {}", repr))
257                    .map_err(|e| PyValueError::new_err(format!("{:?}", e)))?;
258
259                Ok(())
260            }
261            name => Err(PyTypeError::new_err(format!(
262                "cannot operate on {} values",
263                name
264            ))),
265        }
266    }
267
268    #[args(python_exe = "None")]
269    fn oxidize<'p>(&self, py: Python<'p>, python_exe: Option<&PyAny>) -> PyResult<&'p PyTuple> {
270        let python_exe = match python_exe {
271            Some(p) => p,
272            None => {
273                let sys_module = py.import("sys")?;
274                sys_module.getattr("executable")?
275            }
276        };
277        let python_exe = pyobject_to_pathbuf(py, python_exe)?;
278        let temp_dir = PyTempDir::new(py)?;
279        let collector = self.collector.borrow();
280
281        let mut compiler = BytecodeCompiler::new(&python_exe, temp_dir.path()).map_err(|e| {
282            PyValueError::new_err(format!("error constructing bytecode compiler: {:?}", e))
283        })?;
284
285        let prepared: CompiledResourcesCollection = collector
286            .compile_resources(&mut compiler)
287            .context("compiling resources")
288            .map_err(|e| PyValueError::new_err(format!("error oxidizing: {:?}", e)))?;
289
290        let mut resources = Vec::new();
291
292        for resource in prepared.resources.values() {
293            resources.push(resource_to_pyobject(py, resource)?);
294        }
295
296        let mut file_installs = Vec::new();
297
298        for (path, location, executable) in &prepared.extra_files {
299            let path = path_to_pathlib_path(py, path)?;
300            let data = location
301                .resolve_content()
302                .map_err(|e| PyValueError::new_err(e.to_string()))?;
303            let data = PyBytes::new(py, &data);
304            let executable = executable.to_object(py);
305
306            file_installs.push((path, data, executable).to_object(py));
307        }
308
309        Ok(PyTuple::new(
310            py,
311            &[resources.to_object(py), file_installs.to_object(py)],
312        ))
313    }
314}
oxidized_importer/python_resource_collector.rs

oxidized_importer/
python_resource_collector.rs
