oxidized_importer/
lib.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//! oxidized_importer Python extension.
6
7mod conversion;
8#[allow(clippy::needless_option_as_deref)]
9mod importer;
10#[cfg(windows)]
11mod memory_dll;
12mod package_metadata;
13#[allow(clippy::needless_option_as_deref)]
14mod path_entry_finder;
15mod pkg_resources;
16#[allow(clippy::needless_option_as_deref)]
17mod python_resource_collector;
18mod python_resource_types;
19mod python_resources;
20mod resource_reader;
21mod resource_scanning;
22#[cfg(feature = "zipimport")]
23#[allow(clippy::needless_option_as_deref)]
24mod zip_import;
25
26pub use crate::{
27    importer::{
28        install_path_hook, remove_external_importers, replace_meta_path_importers, ImporterState,
29        OxidizedFinder,
30    },
31    python_resource_collector::PyTempDir,
32    python_resources::{PackedResourcesSource, PythonResourcesState},
33};
34
35#[cfg(feature = "zipimport")]
36pub use crate::zip_import::{OxidizedZipFinder, ZipIndex};
37
38use {
39    crate::{
40        path_entry_finder::OxidizedPathEntryFinder,
41        pkg_resources::{register_pkg_resources_with_module, OxidizedPkgResourcesProvider},
42        python_resources::OxidizedResource,
43        resource_reader::OxidizedResourceReader,
44    },
45    pyo3::{
46        exceptions::{PyImportError, PyValueError},
47        ffi as pyffi,
48        prelude::*,
49        AsPyPointer, FromPyPointer,
50    },
51};
52
53/// Name of Python extension module.
54pub const OXIDIZED_IMPORTER_NAME_STR: &str = "oxidized_importer";
55
56/// Null terminated [OXIDIZED_IMPORTER_NAME_STR].
57pub const OXIDIZED_IMPORTER_NAME: &[u8] = b"oxidized_importer\0";
58
59const DOC: &[u8] = b"A highly-performant importer implemented in Rust\0";
60
61static mut MODULE_DEF: pyffi::PyModuleDef = pyffi::PyModuleDef {
62    m_base: pyffi::PyModuleDef_HEAD_INIT,
63    m_name: OXIDIZED_IMPORTER_NAME.as_ptr() as *const _,
64    m_doc: DOC.as_ptr() as *const _,
65    m_size: std::mem::size_of::<ModuleState>() as isize,
66    m_methods: 0 as *mut _,
67    m_slots: 0 as *mut _,
68    m_traverse: None,
69    m_clear: None,
70    m_free: None,
71};
72
73/// State associated with each importer module instance.
74///
75/// We write per-module state to per-module instances of this struct so
76/// we don't rely on global variables and so multiple importer modules can
77/// exist without issue.
78#[derive(Debug)]
79pub(crate) struct ModuleState {
80    /// Whether the module has been initialized.
81    pub(crate) initialized: bool,
82}
83
84/// Obtain the module state for an instance of our importer module.
85///
86/// Creates a Python exception on failure.
87///
88/// Doesn't do type checking that the PyModule is of the appropriate type.
89pub(crate) fn get_module_state(m: &PyModule) -> Result<&mut ModuleState, PyErr> {
90    let ptr = m.as_ptr();
91    let state = unsafe { pyffi::PyModule_GetState(ptr) as *mut ModuleState };
92
93    if state.is_null() {
94        return Err(PyValueError::new_err("unable to retrieve module state"));
95    }
96
97    Ok(unsafe { &mut *state })
98}
99
100/// Module initialization function.
101///
102/// This creates the Python module object.
103///
104/// We don't use the macros in the pyo3 crate because they are somewhat
105/// opinionated about how things should work. e.g. they call
106/// PyEval_InitThreads(), which is undesired. We want total control.
107#[allow(non_snake_case)]
108#[no_mangle]
109pub extern "C" fn PyInit_oxidized_importer() -> *mut pyffi::PyObject {
110    let py = unsafe { Python::assume_gil_acquired() };
111
112    let module = unsafe { pyffi::PyModule_Create(&mut MODULE_DEF) } as *mut pyffi::PyObject;
113
114    if module.is_null() {
115        return module;
116    }
117
118    let module = match unsafe { PyModule::from_owned_ptr_or_err(py, module) } {
119        Ok(m) => m,
120        Err(e) => {
121            e.restore(py);
122            return std::ptr::null_mut();
123        }
124    };
125
126    match module_init(py, module) {
127        Ok(()) => module.into_ptr(),
128        Err(e) => {
129            e.restore(py);
130            std::ptr::null_mut()
131        }
132    }
133}
134
135/// Decodes source bytes into a str.
136///
137/// This is effectively a reimplementation of
138/// importlib._bootstrap_external.decode_source().
139#[pyfunction]
140pub(crate) fn decode_source<'p>(
141    py: Python,
142    io_module: &'p PyModule,
143    source_bytes: &PyAny,
144) -> PyResult<&'p PyAny> {
145    // .py based module, so can't be instantiated until importing mechanism
146    // is bootstrapped.
147    let tokenize_module = py.import("tokenize")?;
148
149    let buffer = io_module.getattr("BytesIO")?.call((source_bytes,), None)?;
150    let readline = buffer.getattr("readline")?;
151    let encoding = tokenize_module
152        .getattr("detect_encoding")?
153        .call((readline,), None)?;
154    let newline_decoder = io_module
155        .getattr("IncrementalNewlineDecoder")?
156        .call((py.None(), true), None)?;
157    let data = source_bytes.call_method("decode", (encoding.get_item(0)?,), None)?;
158    newline_decoder.call_method("decode", (data,), None)
159}
160
161#[pyfunction]
162fn register_pkg_resources(py: Python) -> PyResult<()> {
163    register_pkg_resources_with_module(py, py.import("pkg_resources")?)
164}
165
166/// Initialize the Python module object.
167///
168/// This is called as part of the PyInit_* function to create the internal
169/// module object for the interpreter.
170///
171/// This receives a handle to the current Python interpreter and just-created
172/// Python module instance. It populates the internal module state and registers
173/// functions on the module object for usage by Python.
174fn module_init(py: Python, m: &PyModule) -> PyResult<()> {
175    // Enforce minimum Python version requirement.
176    //
177    // Some features likely work on older Python versions. But we can't
178    // guarantee it. Let's prevent footguns.
179    if py.version_info() < (3, 8) {
180        return Err(PyImportError::new_err("module requires Python 3.8+"));
181    }
182
183    let mut state = get_module_state(m)?;
184
185    state.initialized = false;
186
187    crate::pkg_resources::init_module(m)?;
188    crate::resource_scanning::init_module(m)?;
189
190    m.add_function(wrap_pyfunction!(decode_source, m)?)?;
191    m.add_function(wrap_pyfunction!(register_pkg_resources, m)?)?;
192
193    m.add_class::<crate::package_metadata::OxidizedDistribution>()?;
194    m.add_class::<OxidizedFinder>()?;
195    m.add_class::<OxidizedResource>()?;
196    m.add_class::<crate::python_resource_collector::OxidizedResourceCollector>()?;
197    m.add_class::<OxidizedResourceReader>()?;
198    m.add_class::<OxidizedPathEntryFinder>()?;
199    m.add_class::<OxidizedPkgResourcesProvider>()?;
200    m.add_class::<crate::python_resource_types::PythonModuleSource>()?;
201    m.add_class::<crate::python_resource_types::PythonModuleBytecode>()?;
202    m.add_class::<crate::python_resource_types::PythonPackageResource>()?;
203    m.add_class::<crate::python_resource_types::PythonPackageDistributionResource>()?;
204    m.add_class::<crate::python_resource_types::PythonExtensionModule>()?;
205
206    init_zipimport(m)?;
207
208    Ok(())
209}
210
211#[cfg(feature = "zipimport")]
212fn init_zipimport(m: &PyModule) -> PyResult<()> {
213    m.add_class::<crate::zip_import::OxidizedZipFinder>()?;
214
215    Ok(())
216}
217
218#[cfg(not(feature = "zipimport"))]
219fn init_zipimport(_m: &PyModule) -> PyResult<()> {
220    Ok(())
221}