1use {
10 crate::conversion::{
11 path_to_pathlib_path, pyobject_optional_resources_map_to_owned_bytes,
12 pyobject_optional_resources_map_to_pathbuf, pyobject_to_owned_bytes_optional,
13 pyobject_to_pathbuf_optional,
14 },
15 anyhow::Result,
16 pyo3::{
17 buffer::PyBuffer,
18 exceptions::{PyImportError, PyOSError, PyValueError},
19 ffi as pyffi,
20 prelude::*,
21 types::{PyBytes, PyDict, PyList, PyString, PyTuple},
22 PyTypeInfo,
23 },
24 python_packaging::resource::BytecodeOptimizationLevel,
25 python_packed_resources::Resource,
26 std::{
27 borrow::Cow,
28 cell::RefCell,
29 collections::{hash_map::Entry, BTreeSet, HashMap},
30 ffi::CStr,
31 os::raw::c_int,
32 path::{Path, PathBuf},
33 },
34};
35
36const ENOENT: c_int = 2;
37
38fn is_module_importable<X>(entry: &Resource<X>, optimize_level: BytecodeOptimizationLevel) -> bool
42where
43 [X]: ToOwned<Owned = Vec<X>>,
44{
45 entry.in_memory_source.is_some()
46 || entry.relative_path_module_source.is_some()
47 || match optimize_level {
48 BytecodeOptimizationLevel::Zero => {
49 entry.in_memory_bytecode.is_some() || entry.relative_path_module_bytecode.is_some()
50 }
51 BytecodeOptimizationLevel::One => {
52 entry.in_memory_bytecode_opt1.is_some() || entry.in_memory_bytecode_opt1.is_some()
53 }
54 BytecodeOptimizationLevel::Two => {
55 entry.in_memory_bytecode_opt2.is_some() || entry.in_memory_bytecode_opt2.is_some()
56 }
57 }
58}
59
60pub(crate) fn name_at_package_hierarchy(fullname: &str, package_target: Option<&str>) -> bool {
75 match package_target {
76 None => !fullname.contains('.'),
77 Some(package) => match fullname.strip_prefix(&format!("{}.", package)) {
78 Some(suffix) => !suffix.contains('.'),
79 None => false,
80 },
81 }
82}
83
84pub(crate) fn name_within_package_hierarchy(fullname: &str, package_target: Option<&str>) -> bool {
91 match package_target {
92 None => true,
93 Some(package) => fullname.starts_with(&format!("{}.", package)),
94 }
95}
96
97#[derive(Debug, PartialEq, Eq)]
99pub enum ModuleFlavor {
100 Builtin,
101 Frozen,
102 Extension,
103 SourceBytecode,
104}
105
106pub struct ImportablePythonModule<'a, X: 'a>
111where
112 [X]: ToOwned<Owned = Vec<X>>,
113{
114 resource: &'a Resource<'a, X>,
116
117 current_exe: &'a Path,
119
120 origin: &'a Path,
122
123 pub flavor: ModuleFlavor,
125 pub is_package: bool,
127}
128
129impl<'a> ImportablePythonModule<'a, u8> {
130 pub fn resolve_source<'p>(
136 &self,
137 py: Python<'p>,
138 decode_source: &'p PyAny,
139 io_module: &PyAny,
140 ) -> PyResult<Option<&'p PyAny>> {
141 let bytes = if let Some(data) = &self.resource.in_memory_source {
142 Some(PyBytes::new(py, data))
143 } else if let Some(relative_path) = &self.resource.relative_path_module_source {
144 let path = self.origin.join(relative_path);
145
146 let source = std::fs::read(&path).map_err(|e| {
147 PyErr::from_type(
148 PyImportError::type_object(py),
149 (
150 format!("error reading module source from {}: {}", path.display(), e),
151 self.resource.name.clone().into_py(py),
152 ),
153 )
154 })?;
155
156 Some(PyBytes::new(py, &source))
157 } else {
158 None
159 };
160
161 if let Some(bytes) = bytes {
162 Ok(Some(decode_source.call((io_module, bytes), None)?))
163 } else {
164 Ok(None)
165 }
166 }
167
168 pub fn resolve_bytecode(
177 &mut self,
178 py: Python,
179 optimize_level: BytecodeOptimizationLevel,
180 decode_source: &PyAny,
181 io_module: &PyModule,
182 ) -> PyResult<Option<Py<PyAny>>> {
183 if let Some(data) = match optimize_level {
184 BytecodeOptimizationLevel::Zero => &self.resource.in_memory_bytecode,
185 BytecodeOptimizationLevel::One => &self.resource.in_memory_bytecode_opt1,
186 BytecodeOptimizationLevel::Two => &self.resource.in_memory_bytecode_opt2,
187 } {
188 let ptr = unsafe {
189 pyffi::PyMemoryView_FromMemory(
190 data.as_ptr() as _,
191 data.len() as _,
192 pyffi::PyBUF_READ,
193 )
194 };
195
196 if ptr.is_null() {
197 Ok(None)
198 } else {
199 Ok(Some(unsafe { PyObject::from_owned_ptr(py, ptr) }))
200 }
201 } else if let Some(path) = self.bytecode_path(optimize_level) {
202 let bytecode = std::fs::read(&path).map_err(|e| {
205 PyErr::from_type(
206 PyImportError::type_object(py),
207 (
208 format!("error reading bytecode from {}: {}", path.display(), e)
209 .into_py(py),
210 self.resource.name.clone().into_py(py),
211 ),
212 )
213 })?;
214
215 if bytecode.len() < 16 {
216 return Err(PyImportError::new_err(
217 "bytecode file does not contain enough data",
218 ));
219 }
220
221 Ok(Some(PyBytes::new(py, &bytecode[16..]).into_py(py)))
223 } else if let Some(source) = self.resolve_source(py, decode_source, io_module)? {
224 let builtins = py.import("builtins")?;
225 let marshal = py.import("marshal")?;
226
227 let code = builtins
228 .getattr("compile")?
229 .call((source, self.resource.name.as_ref(), "exec"), None)?;
230 let bytecode = marshal.getattr("dumps")?.call((code,), None)?;
231
232 Ok(Some(bytecode.into_py(py)))
233 } else {
234 Ok(None)
235 }
236 }
237
238 pub fn resolve_module_spec<'p>(
240 &self,
241 py: Python,
242 module_spec_type: &'p PyAny,
243 loader: &PyAny,
244 optimize_level: BytecodeOptimizationLevel,
245 ) -> PyResult<&'p PyAny> {
246 let name = PyString::new(py, &self.resource.name);
247
248 let kwargs = PyDict::new(py);
249 kwargs.set_item("is_package", self.is_package)?;
250
251 let origin = self.resolve_origin(py)?;
256 if let Some(origin) = &origin {
257 kwargs.set_item("origin", origin)?;
258 }
259
260 let spec = module_spec_type.call((name, loader), Some(kwargs))?;
261
262 if origin.is_some() {
263 spec.setattr("has_location", true)?;
264 }
265
266 if let Some(cached) = self.resolve_cached(py, optimize_level)? {
268 spec.setattr("cached", cached)?;
269 }
270
271 if self.is_package {
282 let mut locations = if let Some(origin_path) = self.origin_path() {
300 if let Some(parent_path) = origin_path.parent() {
301 vec![parent_path.into_py(py).into_ref(py)]
302 } else {
303 vec![]
304 }
305 } else {
306 vec![]
307 };
308
309 if locations.is_empty() {
310 let mut path = self.current_exe.to_path_buf();
311 path.extend(self.resource.name.split('.'));
312
313 locations.push(path.into_py(py).into_ref(py));
314 }
315
316 spec.setattr("submodule_search_locations", locations)?;
317 }
318
319 Ok(spec)
320 }
321
322 pub fn resolve_origin<'p>(&self, py: Python<'p>) -> PyResult<Option<&'p PyAny>> {
326 Ok(if let Some(path) = self.origin_path() {
327 Some(path.into_py(py).into_ref(py))
328 } else {
329 None
330 })
331 }
332
333 fn resolve_cached<'p>(
337 &self,
338 py: Python<'p>,
339 optimize_level: BytecodeOptimizationLevel,
340 ) -> PyResult<Option<&'p PyAny>> {
341 let path = match self.flavor {
342 ModuleFlavor::SourceBytecode => self.bytecode_path(optimize_level),
343 _ => None,
344 };
345
346 Ok(if let Some(path) = path {
347 Some(path.into_py(py).into_ref(py))
348 } else {
349 None
350 })
351 }
352
353 fn origin_path(&self) -> Option<PathBuf> {
355 match self.flavor {
356 ModuleFlavor::SourceBytecode => self
357 .resource
358 .relative_path_module_source
359 .as_ref()
360 .map(|path| self.origin.join(path)),
361 ModuleFlavor::Extension => self
362 .resource
363 .relative_path_extension_module_shared_library
364 .as_ref()
365 .map(|path| self.origin.join(path)),
366 _ => None,
367 }
368 }
369
370 fn bytecode_path(&self, optimize_level: BytecodeOptimizationLevel) -> Option<PathBuf> {
372 let bytecode_path = match optimize_level {
373 BytecodeOptimizationLevel::Zero => &self.resource.relative_path_module_bytecode,
374 BytecodeOptimizationLevel::One => &self.resource.relative_path_module_bytecode_opt1,
375 BytecodeOptimizationLevel::Two => &self.resource.relative_path_module_bytecode_opt2,
376 };
377
378 bytecode_path
379 .as_ref()
380 .map(|bytecode_path| self.origin.join(bytecode_path))
381 }
382
383 pub fn in_memory_extension_module_shared_library(&self) -> &'a Option<Cow<'a, [u8]>> {
384 &self.resource.in_memory_extension_module_shared_library
385 }
386}
387
388#[derive(Clone, Debug, PartialEq, Eq)]
390pub enum PackedResourcesSource<'a> {
391 Memory(&'a [u8]),
393
394 #[allow(unused)]
396 MemoryMappedPath(PathBuf),
397}
398
399impl<'a> From<&'a [u8]> for PackedResourcesSource<'a> {
400 fn from(data: &'a [u8]) -> Self {
401 Self::Memory(data)
402 }
403}
404
405#[derive(Debug)]
407pub struct PythonResourcesState<'a, X>
408where
409 [X]: ToOwned<Owned = Vec<X>>,
410{
411 current_exe: PathBuf,
413
414 origin: PathBuf,
418
419 resources: HashMap<Cow<'a, str>, Resource<'a, X>>,
421
422 backing_py_objects: Vec<Py<PyAny>>,
427
428 backing_mmaps: Vec<memmap2::Mmap>,
430}
431
432impl<'a> Default for PythonResourcesState<'a, u8> {
433 fn default() -> Self {
434 Self {
435 current_exe: PathBuf::new(),
436 origin: PathBuf::new(),
437 resources: HashMap::new(),
438 backing_py_objects: vec![],
439 backing_mmaps: vec![],
440 }
441 }
442}
443
444impl<'a> PythonResourcesState<'a, u8> {
445 pub fn new_from_env() -> Result<Self, &'static str> {
447 let exe = std::env::current_exe().map_err(|_| "unable to obtain current executable")?;
448 let origin = exe
449 .parent()
450 .ok_or("unable to get executable parent")?
451 .to_path_buf();
452
453 Ok(Self {
454 current_exe: exe,
455 origin,
456 ..Default::default()
457 })
458 }
459
460 pub fn current_exe(&self) -> &Path {
462 &self.current_exe
463 }
464
465 pub fn set_current_exe(&mut self, path: PathBuf) {
467 self.current_exe = path;
468 }
469
470 pub fn origin(&self) -> &Path {
472 &self.origin
473 }
474
475 pub fn set_origin(&mut self, path: PathBuf) {
477 self.origin = path;
478 }
479
480 pub fn index_data(&mut self, data: &'a [u8]) -> Result<(), &'static str> {
487 let resources = python_packed_resources::load_resources(data)?;
488
489 self.resources.reserve(resources.expected_resources_count());
492
493 for resource in resources {
494 let resource = resource?;
495
496 match self.resources.entry(resource.name.clone()) {
497 Entry::Occupied(existing) => {
498 existing.into_mut().merge_from(resource)?;
499 }
500 Entry::Vacant(vacant) => {
501 vacant.insert(resource);
502 }
503 }
504 }
505
506 Ok(())
507 }
508
509 pub fn index_path_memory_mapped(&mut self, path: impl AsRef<Path>) -> Result<(), String> {
511 let path = path.as_ref();
512 let f = std::fs::File::open(path).map_err(|e| e.to_string())?;
513
514 let mapped = unsafe { memmap2::Mmap::map(&f) }.map_err(|e| e.to_string())?;
515
516 let data = unsafe { std::slice::from_raw_parts::<u8>(mapped.as_ptr(), mapped.len()) };
517
518 self.index_data(data)?;
519 self.backing_mmaps.push(mapped);
520
521 Ok(())
522 }
523
524 pub fn index_pyobject(&mut self, py: Python, obj: &PyAny) -> PyResult<()> {
528 let buffer = PyBuffer::<u8>::get(obj)?;
529
530 let data = unsafe {
531 std::slice::from_raw_parts::<u8>(buffer.buf_ptr() as *const _, buffer.len_bytes())
532 };
533
534 self.index_data(data).map_err(PyValueError::new_err)?;
535 self.backing_py_objects.push(obj.to_object(py));
536
537 Ok(())
538 }
539
540 pub fn index_interpreter_builtin_extension_modules(&mut self) -> Result<(), &'static str> {
542 for i in 0.. {
543 let record = unsafe { pyffi::PyImport_Inittab.offset(i) };
544
545 if unsafe { *record }.name.is_null() {
546 break;
547 }
548
549 let name = unsafe { CStr::from_ptr((*record).name as _) };
550 let name_str = match name.to_str() {
551 Ok(v) => v,
552 Err(_) => {
553 return Err("unable to parse PyImport_Inittab");
554 }
555 };
556
557 self.resources
558 .entry(name_str.into())
559 .and_modify(|r| {
560 r.is_python_builtin_extension_module = true;
561 })
562 .or_insert_with(|| Resource {
563 is_python_builtin_extension_module: true,
564 name: Cow::Owned(name_str.to_string()),
565 ..Resource::default()
566 });
567 }
568
569 Ok(())
570 }
571
572 pub fn index_interpreter_frozen_modules(&mut self) -> Result<(), &'static str> {
574 for i in 0.. {
575 let record = unsafe { pyffi::PyImport_FrozenModules.offset(i) };
576
577 if unsafe { *record }.name.is_null() {
578 break;
579 }
580
581 let name = unsafe { CStr::from_ptr((*record).name as _) };
582 let name_str = match name.to_str() {
583 Ok(v) => v,
584 Err(_) => {
585 return Err("unable to parse PyImport_FrozenModules");
586 }
587 };
588
589 self.resources
590 .entry(name_str.into())
591 .and_modify(|r| {
592 r.is_python_frozen_module = true;
593 })
594 .or_insert_with(|| Resource {
595 is_python_frozen_module: true,
596 name: Cow::Owned(name_str.to_string()),
597 ..Resource::default()
598 });
599 }
600
601 Ok(())
602 }
603
604 pub fn index_interpreter_builtins(&mut self) -> Result<(), &'static str> {
610 self.index_interpreter_builtin_extension_modules()?;
611 self.index_interpreter_frozen_modules()?;
612
613 Ok(())
614 }
615
616 pub fn has_resource(&self, name: &str) -> bool {
618 self.resources.contains_key(name)
619 }
620
621 pub fn add_resource<'resource: 'a>(
626 &mut self,
627 resource: Resource<'resource, u8>,
628 ) -> Result<(), &'static str> {
629 self.resources.insert(resource.name.clone(), resource);
630
631 Ok(())
632 }
633
634 pub fn resolve_importable_module(
636 &self,
637 name: &str,
638 optimize_level: BytecodeOptimizationLevel,
639 ) -> Option<ImportablePythonModule<u8>> {
640 let name = name.strip_suffix(".__init__").unwrap_or(name);
664
665 let resource = match self.resources.get(name) {
666 Some(entry) => entry,
667 None => return None,
668 };
669
670 if resource.is_python_builtin_extension_module {
697 Some(ImportablePythonModule {
698 resource,
699 current_exe: &self.current_exe,
700 origin: &self.origin,
701 flavor: ModuleFlavor::Builtin,
702 is_package: resource.is_python_package,
703 })
704 } else if resource.is_python_frozen_module {
705 Some(ImportablePythonModule {
706 resource,
707 current_exe: &self.current_exe,
708 origin: &self.origin,
709 flavor: ModuleFlavor::Frozen,
710 is_package: resource.is_python_package,
711 })
712 } else if resource.is_python_extension_module {
713 Some(ImportablePythonModule {
714 resource,
715 current_exe: &self.current_exe,
716 origin: &self.origin,
717 flavor: ModuleFlavor::Extension,
718 is_package: resource.is_python_package,
719 })
720 } else if resource.is_python_module {
721 if is_module_importable(resource, optimize_level) {
722 Some(ImportablePythonModule {
723 resource,
724 current_exe: &self.current_exe,
725 origin: &self.origin,
726 flavor: ModuleFlavor::SourceBytecode,
727 is_package: resource.is_python_package,
728 })
729 } else {
730 None
731 }
732 } else {
733 None
734 }
735 }
736
737 pub fn get_package_resource_file<'p>(
743 &self,
744 py: Python<'p>,
745 package: &str,
746 resource_name: &str,
747 ) -> PyResult<Option<&'p PyAny>> {
748 let entry = match self.resources.get(package) {
749 Some(entry) => entry,
750 None => return Ok(None),
751 };
752
753 if let Some(resources) = &entry.in_memory_package_resources {
754 if let Some(data) = resources.get(resource_name) {
755 let io_module = py.import("io")?;
756 let bytes_io = io_module.getattr("BytesIO")?;
757
758 let data = PyBytes::new(py, data);
759 return Ok(Some(bytes_io.call((data,), None)?));
760 }
761 }
762
763 if let Some(resources) = &entry.relative_path_package_resources {
764 if let Some(path) = resources.get(resource_name) {
765 let path = self.origin.join(path);
766 let io_module = py.import("io")?;
767
768 return Ok(Some(
769 io_module
770 .getattr("FileIO")?
771 .call((path.into_py(py), "r"), None)?,
772 ));
773 }
774 }
775
776 Ok(None)
777 }
778
779 pub fn is_package_resource(&self, package: &str, resource_name: &str) -> bool {
781 if let Some(entry) = self.resources.get(package) {
782 if let Some(resources) = &entry.in_memory_package_resources {
783 if resources.contains_key(resource_name) {
784 return true;
785 }
786 }
787
788 if let Some(resources) = &entry.relative_path_package_resources {
789 if resources.contains_key(resource_name) {
790 return true;
791 }
792 }
793 }
794
795 false
796 }
797
798 pub fn package_resource_names<'p>(&self, py: Python<'p>, package: &str) -> PyResult<&'p PyAny> {
802 let entry = match self.resources.get(package) {
803 Some(entry) => entry,
804 None => return Ok(PyList::empty(py).into()),
805 };
806
807 let mut names = if let Some(resources) = &entry.in_memory_package_resources {
808 resources.keys().collect()
809 } else if let Some(resources) = &entry.relative_path_package_resources {
810 resources.keys().collect()
811 } else {
812 vec![]
813 };
814
815 names.sort();
816
817 let names = names
818 .iter()
819 .map(|x| x.to_object(py))
820 .collect::<Vec<Py<PyAny>>>();
821
822 Ok(PyList::new(py, &names).into())
823 }
824
825 pub fn is_package_resource_directory(&self, package: &str, name: &str) -> bool {
827 let name = name.replace('\\', "/");
829
830 let prefix = if name.ends_with('/') {
831 name
832 } else {
833 format!("{}/", name)
834 };
835
836 if let Some(entry) = self.resources.get(package) {
837 if let Some(resources) = &entry.in_memory_package_resources {
838 if resources.keys().any(|path| path.starts_with(&prefix)) {
839 return true;
840 }
841 }
842
843 if let Some(resources) = &entry.relative_path_package_resources {
844 if resources.keys().any(|path| path.starts_with(&prefix)) {
845 return true;
846 }
847 }
848
849 false
850 } else {
851 false
852 }
853 }
854
855 pub fn package_resources_list_directory(&self, package: &str, name: &str) -> Vec<String> {
857 let name = name.replace('\\', "/");
858
859 let prefix = if name.ends_with('/') {
860 Some(name)
861 } else if name.is_empty() {
862 None
863 } else {
864 Some(format!("{}/", name))
865 };
866
867 let filter_map_resource = |path: &'_ Cow<'_, str>| -> Option<String> {
868 match &prefix {
869 Some(prefix) => {
870 if let Some(name) = path.strip_prefix(prefix) {
871 if name.contains('/') {
872 None
873 } else {
874 Some(name.to_string())
875 }
876 } else {
877 None
878 }
879 }
880 None => {
881 if path.contains('/') {
883 None
884 } else {
885 Some(path.to_string())
886 }
887 }
888 }
889 };
890
891 let mut entries = BTreeSet::new();
892
893 if let Some(entry) = self.resources.get(package) {
894 if let Some(resources) = &entry.in_memory_package_resources {
895 entries.extend(resources.keys().filter_map(filter_map_resource));
896 }
897
898 if let Some(resources) = &entry.relative_path_package_resources {
899 entries.extend(resources.keys().filter_map(filter_map_resource));
900 }
901 }
902
903 entries.into_iter().collect::<Vec<_>>()
904 }
905
906 pub fn resolve_resource_data_from_path<'p>(
913 &self,
914 py: Python<'p>,
915 path: &str,
916 ) -> PyResult<&'p PyAny> {
917 let path = path.to_owned();
953 let native_path = PathBuf::from(&path);
954
955 let (relative_path, check_in_memory, check_relative_path) =
956 if let Ok(relative_path) = native_path.strip_prefix(&self.current_exe) {
957 (relative_path, true, false)
958 } else if let Ok(relative_path) = native_path.strip_prefix(&self.origin) {
959 (relative_path, false, true)
960 } else {
961 return Err(PyErr::from_type(
962 PyOSError::type_object(py),
963 (ENOENT, "resource not known", path),
964 ));
965 };
966
967 let components = relative_path.components().collect::<Vec<_>>();
981
982 if components.len() < 2 {
985 return Err(PyErr::from_type(
986 PyOSError::type_object(py),
987 (
988 ENOENT,
989 "illegal resource name: missing package component",
990 path,
991 ),
992 ));
993 }
994
995 let mut name_parts = vec![components[components.len() - 1]
996 .as_os_str()
997 .to_string_lossy()];
998 let mut package_parts = components[0..components.len() - 1]
999 .iter()
1000 .map(|c| c.as_os_str().to_string_lossy())
1001 .collect::<Vec<_>>();
1002
1003 while !package_parts.is_empty() {
1004 let package_name = package_parts.join(".");
1005 let package_name_ref: &str = &package_name;
1006
1007 let resource_name = name_parts.join("/");
1009 let resource_name_ref: &str = &resource_name;
1010
1011 if let Some(entry) = self.resources.get(package_name_ref) {
1012 if check_in_memory {
1013 if let Some(resources) = &entry.in_memory_package_resources {
1014 if let Some(data) = resources.get(resource_name_ref) {
1015 return Ok(PyBytes::new(py, data).into());
1016 }
1017 }
1018 }
1019
1020 if check_relative_path {
1021 if let Some(resources) = &entry.relative_path_package_resources {
1022 if let Some(resource_relative_path) = resources.get(resource_name_ref) {
1023 let resource_path = self.origin.join(resource_relative_path);
1024
1025 let io_module = py.import("io")?;
1026
1027 let fh = io_module
1028 .getattr("FileIO")?
1029 .call((resource_path.into_py(py).into_ref(py), "r"), None)?;
1030
1031 return fh.call_method0("read");
1032 }
1033 }
1034 }
1035
1036 break;
1039 }
1040
1041 name_parts.insert(0, package_parts.pop().unwrap());
1042 }
1043
1044 Err(PyErr::from_type(
1047 PyOSError::type_object(py),
1048 (ENOENT, "resource not known", path),
1049 ))
1050 }
1051
1052 pub fn pkgutil_modules_infos<'p>(
1059 &self,
1060 py: Python<'p>,
1061 package_filter: Option<&str>,
1062 prefix: Option<String>,
1063 optimize_level: BytecodeOptimizationLevel,
1064 ) -> PyResult<&'p PyList> {
1065 let infos: PyResult<Vec<_>> = self
1066 .resources
1067 .values()
1068 .filter(|r| {
1069 r.is_python_extension_module
1070 || (r.is_python_module && is_module_importable(r, optimize_level))
1071 })
1072 .filter(|r| name_at_package_hierarchy(&r.name, package_filter))
1073 .map(|r| {
1074 let name = r.name.rsplit('.').next().unwrap();
1076
1077 let name = if let Some(prefix) = &prefix {
1078 format!("{}{}", prefix, name)
1079 } else {
1080 name.to_string()
1081 };
1082
1083 let name = name.to_object(py);
1084 let is_package = r.is_python_package.to_object(py);
1085
1086 Ok(PyTuple::new(py, &[name, is_package]))
1087 })
1088 .collect();
1089
1090 let infos = infos?;
1091
1092 Ok(PyList::new(py, &infos))
1093 }
1094
1095 pub fn package_distribution_names(&self, filter: impl Fn(&str) -> bool) -> Vec<&'_ str> {
1097 self.resources
1098 .values()
1099 .filter(|r| {
1100 r.is_python_package
1101 && (r.in_memory_distribution_resources.is_some()
1102 || r.relative_path_distribution_resources.is_some())
1103 })
1104 .filter(|r| filter(r.name.as_ref()))
1105 .map(|r| r.name.as_ref())
1106 .collect::<Vec<_>>()
1107 }
1108
1109 pub fn resolve_package_distribution_resource(
1111 &self,
1112 package: &str,
1113 name: &str,
1114 ) -> Result<Option<Cow<'_, [u8]>>> {
1115 if let Some(entry) = self.resources.get(package) {
1116 if let Some(resources) = &entry.in_memory_distribution_resources {
1117 if let Some(data) = resources.get(name) {
1118 return Ok(Some(Cow::Borrowed(data.as_ref())));
1119 }
1120 }
1121
1122 if let Some(resources) = &entry.relative_path_distribution_resources {
1123 if let Some(path) = resources.get(name) {
1124 let path = &self.origin.join(path);
1125 let data = std::fs::read(path)?;
1126
1127 return Ok(Some(Cow::Owned(data)));
1128 }
1129 }
1130
1131 Ok(None)
1132 } else {
1133 Ok(None)
1134 }
1135 }
1136
1137 pub fn package_distribution_resource_name_is_directory(
1139 &self,
1140 package: &str,
1141 name: &str,
1142 ) -> bool {
1143 let name = name.replace('\\', "/");
1144
1145 let prefix = if name.ends_with('/') {
1146 name
1147 } else {
1148 format!("{}/", name)
1149 };
1150
1151 if let Some(entry) = &self.resources.get(package) {
1152 if let Some(resources) = &entry.in_memory_distribution_resources {
1153 if resources.keys().any(|path| path.starts_with(&prefix)) {
1154 return true;
1155 }
1156 }
1157
1158 if let Some(resources) = &entry.relative_path_distribution_resources {
1159 if resources.keys().any(|path| path.starts_with(&prefix)) {
1160 return true;
1161 }
1162 }
1163
1164 false
1165 } else {
1166 false
1167 }
1168 }
1169
1170 pub fn package_distribution_resources_list_directory<'slf>(
1172 &'slf self,
1173 package: &str,
1174 name: &str,
1175 ) -> Vec<&'slf str> {
1176 let name = name.replace('\\', "/");
1177
1178 let prefix = if name.ends_with('/') {
1179 Some(name)
1180 } else if name.is_empty() {
1181 None
1182 } else {
1183 Some(format!("{}/", name))
1184 };
1185
1186 let filter_map_resource = |path: &'slf Cow<'slf, str>| -> Option<&'slf str> {
1187 match &prefix {
1188 Some(prefix) => {
1189 path.strip_prefix(prefix).filter(|&name| !name.contains('/'))
1190 }
1191 None => {
1192 if path.contains('/') {
1194 None
1195 } else {
1196 Some(path)
1197 }
1198 }
1199 }
1200 };
1201
1202 let mut entries = BTreeSet::new();
1203
1204 if let Some(entry) = self.resources.get(package) {
1205 if let Some(resources) = &entry.in_memory_distribution_resources {
1206 entries.extend(resources.keys().filter_map(filter_map_resource));
1207 }
1208
1209 if let Some(resources) = &entry.relative_path_distribution_resources {
1210 entries.extend(resources.keys().filter_map(filter_map_resource));
1211 }
1212 }
1213
1214 entries.into_iter().collect::<Vec<_>>()
1215 }
1216
1217 pub fn resolve_in_memory_shared_library_data(&self, name: &str) -> Option<&[u8]> {
1219 if let Some(entry) = &self.resources.get(name) {
1220 if let Some(library_data) = &entry.in_memory_shared_library {
1221 Some(library_data.as_ref())
1222 } else {
1223 None
1224 }
1225 } else {
1226 None
1227 }
1228 }
1229
1230 pub fn resources_as_py_list<'p>(&self, py: Python<'p>) -> PyResult<&'p PyList> {
1232 let mut resources = self.resources.values().collect::<Vec<_>>();
1233 resources.sort_by_key(|r| &r.name);
1234
1235 let objects = resources
1236 .iter()
1237 .map(|r| resource_to_pyobject(py, r))
1238 .collect::<Result<Vec<_>, _>>()?;
1239
1240 Ok(PyList::new(py, objects))
1241 }
1242
1243 pub fn serialize_resources(
1248 &self,
1249 ignore_builtin: bool,
1250 ignore_frozen: bool,
1251 ) -> Result<Vec<u8>> {
1252 let mut resources = self
1253 .resources
1254 .values()
1255 .filter(|resource| {
1256 !((resource.is_python_builtin_extension_module && ignore_builtin)
1258 || (resource.is_python_frozen_module && ignore_frozen))
1259 })
1260 .collect::<Vec<&Resource<u8>>>();
1261
1262 resources.sort_by_key(|v| &v.name);
1264
1265 let mut buffer = Vec::new();
1266
1267 python_packed_resources::write_packed_resources_v3(&resources, &mut buffer, None)?;
1268
1269 Ok(buffer)
1270 }
1271}
1272
1273#[pyclass(module = "oxidized_importer")]
1274pub(crate) struct OxidizedResource {
1275 resource: RefCell<Resource<'static, u8>>,
1276}
1277
1278#[pymethods]
1279impl OxidizedResource {
1280 fn __repr__(&self) -> String {
1281 format!(
1282 "<OxidizedResource name=\"{}\">",
1283 self.resource.borrow().name
1284 )
1285 }
1286
1287 #[new]
1288 fn new() -> PyResult<Self> {
1289 Ok(Self {
1290 resource: RefCell::new(Resource::<u8>::default()),
1291 })
1292 }
1293
1294 #[getter]
1295 fn get_is_module(&self) -> bool {
1296 self.resource.borrow().is_python_module
1297 }
1298
1299 #[setter]
1300 fn set_is_module(&self, value: bool) -> PyResult<()> {
1301 self.resource.borrow_mut().is_python_module = value;
1302
1303 Ok(())
1304 }
1305
1306 #[getter]
1307 fn get_is_builtin_extension_module(&self) -> bool {
1308 self.resource.borrow().is_python_builtin_extension_module
1309 }
1310
1311 #[setter]
1312 fn set_is_builtin_extension_module(&self, value: bool) -> PyResult<()> {
1313 self.resource
1314 .borrow_mut()
1315 .is_python_builtin_extension_module = value;
1316
1317 Ok(())
1318 }
1319
1320 #[getter]
1321 fn get_is_frozen_module(&self) -> bool {
1322 self.resource.borrow().is_python_frozen_module
1323 }
1324
1325 #[setter]
1326 fn set_is_frozen_module(&self, value: bool) -> PyResult<()> {
1327 self.resource.borrow_mut().is_python_frozen_module = value;
1328
1329 Ok(())
1330 }
1331
1332 #[getter]
1333 fn get_is_extension_module(&self) -> bool {
1334 self.resource.borrow().is_python_extension_module
1335 }
1336
1337 #[setter]
1338 fn set_is_extension_module(&self, value: bool) -> PyResult<()> {
1339 self.resource.borrow_mut().is_python_extension_module = value;
1340
1341 Ok(())
1342 }
1343
1344 #[getter]
1345 fn get_is_shared_library(&self) -> bool {
1346 self.resource.borrow().is_shared_library
1347 }
1348
1349 #[setter]
1350 fn set_is_shared_library(&self, value: bool) -> PyResult<()> {
1351 self.resource.borrow_mut().is_shared_library = value;
1352
1353 Ok(())
1354 }
1355
1356 #[getter]
1357 fn get_name(&self) -> String {
1358 self.resource.borrow().name.to_string()
1359 }
1360
1361 #[setter]
1362 fn set_name(&self, value: &str) -> PyResult<()> {
1363 self.resource.borrow_mut().name = Cow::Owned(value.to_owned());
1364
1365 Ok(())
1366 }
1367
1368 #[getter]
1369 fn get_is_package(&self) -> bool {
1370 self.resource.borrow().is_python_package
1371 }
1372
1373 #[setter]
1374 fn set_is_package(&self, value: bool) -> PyResult<()> {
1375 self.resource.borrow_mut().is_python_package = value;
1376
1377 Ok(())
1378 }
1379
1380 #[getter]
1381 fn get_is_namespace_package(&self) -> bool {
1382 self.resource.borrow().is_python_namespace_package
1383 }
1384
1385 #[setter]
1386 fn set_is_namespace_package(&self, value: bool) -> PyResult<()> {
1387 self.resource.borrow_mut().is_python_namespace_package = value;
1388
1389 Ok(())
1390 }
1391
1392 #[getter]
1393 fn get_in_memory_source<'p>(&self, py: Python<'p>) -> Option<&'p PyBytes> {
1394 self.resource
1395 .borrow()
1396 .in_memory_source
1397 .as_ref()
1398 .map(|x| PyBytes::new(py, x))
1399 }
1400
1401 #[setter]
1402 fn set_in_memory_source(&self, value: &PyAny) -> PyResult<()> {
1403 self.resource.borrow_mut().in_memory_source =
1404 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1405
1406 Ok(())
1407 }
1408
1409 #[getter]
1410 fn get_in_memory_bytecode<'p>(&self, py: Python<'p>) -> Option<&'p PyBytes> {
1411 self.resource
1412 .borrow()
1413 .in_memory_bytecode
1414 .as_ref()
1415 .map(|x| PyBytes::new(py, x))
1416 }
1417
1418 #[setter]
1419 fn set_in_memory_bytecode(&self, value: &PyAny) -> PyResult<()> {
1420 self.resource.borrow_mut().in_memory_bytecode =
1421 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1422
1423 Ok(())
1424 }
1425
1426 #[getter]
1427 fn get_in_memory_bytecode_opt1<'p>(&self, py: Python<'p>) -> Option<&'p PyBytes> {
1428 self.resource
1429 .borrow()
1430 .in_memory_bytecode_opt1
1431 .as_ref()
1432 .map(|x| PyBytes::new(py, x))
1433 }
1434
1435 #[setter]
1436 fn set_in_memory_bytecode_opt1(&self, value: &PyAny) -> PyResult<()> {
1437 self.resource.borrow_mut().in_memory_bytecode_opt1 =
1438 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1439
1440 Ok(())
1441 }
1442
1443 #[getter]
1444 fn get_in_memory_bytecode_opt2<'p>(&self, py: Python<'p>) -> Option<&'p PyBytes> {
1445 self.resource
1446 .borrow()
1447 .in_memory_bytecode_opt2
1448 .as_ref()
1449 .map(|x| PyBytes::new(py, x))
1450 }
1451
1452 #[setter]
1453 fn set_in_memory_bytecode_opt2(&self, value: &PyAny) -> PyResult<()> {
1454 self.resource.borrow_mut().in_memory_bytecode_opt2 =
1455 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1456
1457 Ok(())
1458 }
1459
1460 #[getter]
1461 fn get_in_memory_extension_module_shared_library<'p>(
1462 &self,
1463 py: Python<'p>,
1464 ) -> Option<&'p PyBytes> {
1465 self.resource
1466 .borrow()
1467 .in_memory_extension_module_shared_library
1468 .as_ref()
1469 .map(|x| PyBytes::new(py, x))
1470 }
1471
1472 #[setter]
1473 fn set_in_memory_extension_module_shared_library(&self, value: &PyAny) -> PyResult<()> {
1474 self.resource
1475 .borrow_mut()
1476 .in_memory_extension_module_shared_library =
1477 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1478
1479 Ok(())
1480 }
1481
1482 #[getter]
1483 fn get_in_memory_package_resources<'p>(
1484 &self,
1485 py: Python<'p>,
1486 ) -> Option<HashMap<String, &'p PyBytes>> {
1487 self.resource
1488 .borrow()
1489 .in_memory_package_resources
1490 .as_ref()
1491 .map(|x| {
1492 x.iter()
1493 .map(|(k, v)| (k.to_string(), PyBytes::new(py, v)))
1494 .collect()
1495 })
1496 }
1497
1498 #[setter]
1499 fn set_in_memory_package_resources(&self, value: &PyAny) -> PyResult<()> {
1500 self.resource.borrow_mut().in_memory_package_resources =
1501 pyobject_optional_resources_map_to_owned_bytes(value)?.map(|x| {
1502 x.into_iter()
1503 .map(|(k, v)| (Cow::Owned(k), Cow::Owned(v)))
1504 .collect()
1505 });
1506
1507 Ok(())
1508 }
1509
1510 #[getter]
1511 fn get_in_memory_distribution_resources<'p>(
1512 &self,
1513 py: Python<'p>,
1514 ) -> Option<HashMap<String, &'p PyBytes>> {
1515 self.resource
1516 .borrow()
1517 .in_memory_distribution_resources
1518 .as_ref()
1519 .map(|x| {
1520 x.iter()
1521 .map(|(k, v)| (k.to_string(), PyBytes::new(py, v)))
1522 .collect()
1523 })
1524 }
1525
1526 #[setter]
1527 fn set_in_memory_distribution_resources(&self, value: &PyAny) -> PyResult<()> {
1528 self.resource.borrow_mut().in_memory_distribution_resources =
1529 pyobject_optional_resources_map_to_owned_bytes(value)?.map(|x| {
1530 x.into_iter()
1531 .map(|(k, v)| (Cow::Owned(k), Cow::Owned(v)))
1532 .collect()
1533 });
1534
1535 Ok(())
1536 }
1537
1538 #[getter]
1539 fn get_in_memory_shared_library<'p>(&self, py: Python<'p>) -> Option<&'p PyBytes> {
1540 self.resource
1541 .borrow()
1542 .in_memory_shared_library
1543 .as_ref()
1544 .map(|x| PyBytes::new(py, x))
1545 }
1546
1547 #[setter]
1548 fn set_in_memory_shared_library(&self, value: &PyAny) -> PyResult<()> {
1549 self.resource.borrow_mut().in_memory_shared_library =
1550 pyobject_to_owned_bytes_optional(value)?.map(Cow::Owned);
1551
1552 Ok(())
1553 }
1554
1555 #[getter]
1556 fn get_shared_library_dependency_names(&self) -> Option<Vec<String>> {
1557 self.resource
1558 .borrow()
1559 .shared_library_dependency_names
1560 .as_ref()
1561 .map(|x| x.iter().map(|v| v.to_string()).collect())
1562 }
1563
1564 #[setter]
1565 fn set_shared_library_dependency_names(&self, value: Option<Vec<String>>) -> PyResult<()> {
1566 self.resource.borrow_mut().shared_library_dependency_names =
1567 value.map(|x| x.into_iter().map(Cow::Owned).collect());
1568
1569 Ok(())
1570 }
1571
1572 #[getter]
1573 fn get_relative_path_module_source<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1574 self.resource
1575 .borrow()
1576 .relative_path_module_source
1577 .as_ref()
1578 .map_or_else(
1579 || Ok(py.None().into_ref(py)),
1580 |x| path_to_pathlib_path(py, x),
1581 )
1582 }
1583
1584 #[setter]
1585 fn set_relative_path_module_source(&self, py: Python, value: &PyAny) -> PyResult<()> {
1586 self.resource.borrow_mut().relative_path_module_source =
1587 pyobject_to_pathbuf_optional(py, value)?.map(Cow::Owned);
1588
1589 Ok(())
1590 }
1591
1592 #[getter]
1593 fn get_relative_path_module_bytecode<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1594 self.resource
1595 .borrow()
1596 .relative_path_module_bytecode
1597 .as_ref()
1598 .map_or_else(
1599 || Ok(py.None().into_ref(py)),
1600 |x| path_to_pathlib_path(py, x),
1601 )
1602 }
1603
1604 #[setter]
1605 fn set_relative_path_module_bytecode(&self, py: Python, value: &PyAny) -> PyResult<()> {
1606 self.resource.borrow_mut().relative_path_module_bytecode =
1607 pyobject_to_pathbuf_optional(py, value)?.map(Cow::Owned);
1608
1609 Ok(())
1610 }
1611
1612 #[getter]
1613 fn get_relative_path_module_bytecode_opt1<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1614 self.resource
1615 .borrow()
1616 .relative_path_module_bytecode_opt1
1617 .as_ref()
1618 .map_or_else(
1619 || Ok(py.None().into_ref(py)),
1620 |x| path_to_pathlib_path(py, x),
1621 )
1622 }
1623
1624 #[setter]
1625 fn set_relative_path_module_bytecode_opt1(&self, py: Python, value: &PyAny) -> PyResult<()> {
1626 self.resource
1627 .borrow_mut()
1628 .relative_path_module_bytecode_opt1 =
1629 pyobject_to_pathbuf_optional(py, value)?.map(Cow::Owned);
1630
1631 Ok(())
1632 }
1633
1634 #[getter]
1635 fn get_relative_path_module_bytecode_opt2<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1636 self.resource
1637 .borrow()
1638 .relative_path_module_bytecode_opt2
1639 .as_ref()
1640 .map_or_else(
1641 || Ok(py.None().into_ref(py)),
1642 |x| path_to_pathlib_path(py, x),
1643 )
1644 }
1645
1646 #[setter]
1647 fn set_relative_path_module_bytecode_opt2(&self, py: Python, value: &PyAny) -> PyResult<()> {
1648 self.resource
1649 .borrow_mut()
1650 .relative_path_module_bytecode_opt2 =
1651 pyobject_to_pathbuf_optional(py, value)?.map(Cow::Owned);
1652
1653 Ok(())
1654 }
1655
1656 #[getter]
1657 fn get_relative_path_extension_module_shared_library<'p>(
1658 &self,
1659 py: Python<'p>,
1660 ) -> PyResult<&'p PyAny> {
1661 self.resource
1662 .borrow()
1663 .relative_path_extension_module_shared_library
1664 .as_ref()
1665 .map_or_else(
1666 || Ok(py.None().into_ref(py)),
1667 |x| path_to_pathlib_path(py, x),
1668 )
1669 }
1670
1671 #[setter]
1672 fn set_relative_path_extension_module_shared_library(
1673 &self,
1674 py: Python,
1675 value: &PyAny,
1676 ) -> PyResult<()> {
1677 self.resource
1678 .borrow_mut()
1679 .relative_path_extension_module_shared_library =
1680 pyobject_to_pathbuf_optional(py, value)?.map(Cow::Owned);
1681
1682 Ok(())
1683 }
1684
1685 #[getter]
1686 fn get_relative_path_package_resources<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1687 self.resource
1688 .borrow()
1689 .relative_path_package_resources
1690 .as_ref()
1691 .map_or_else(
1692 || Ok(py.None().into_ref(py)),
1693 |x| -> PyResult<&PyAny> {
1694 let res = PyDict::new(py);
1695
1696 for (k, v) in x.iter() {
1697 res.set_item(k, path_to_pathlib_path(py, v)?)?;
1698 }
1699
1700 Ok(res)
1701 },
1702 )
1703 }
1704
1705 #[setter]
1706 fn set_relative_path_package_resources(&self, py: Python, value: &PyAny) -> PyResult<()> {
1707 self.resource.borrow_mut().relative_path_package_resources =
1708 pyobject_optional_resources_map_to_pathbuf(py, value)?.map(|x| {
1709 x.into_iter()
1710 .map(|(k, v)| (Cow::Owned(k), Cow::Owned(v)))
1711 .collect()
1712 });
1713
1714 Ok(())
1715 }
1716
1717 #[getter]
1718 fn get_relative_path_distribution_resources<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
1719 self.resource
1720 .borrow()
1721 .relative_path_distribution_resources
1722 .as_ref()
1723 .map_or_else(
1724 || Ok(py.None().into_ref(py)),
1725 |x| -> PyResult<&PyAny> {
1726 let res = PyDict::new(py);
1727
1728 for (k, v) in x.iter() {
1729 res.set_item(k, path_to_pathlib_path(py, v)?)?;
1730 }
1731
1732 Ok(res.into())
1733 },
1734 )
1735 }
1736
1737 #[setter]
1738 fn set_relative_path_distribution_resources(&self, py: Python, value: &PyAny) -> PyResult<()> {
1739 self.resource
1740 .borrow_mut()
1741 .relative_path_distribution_resources =
1742 pyobject_optional_resources_map_to_pathbuf(py, value)?.map(|x| {
1743 x.into_iter()
1744 .map(|(k, v)| (Cow::Owned(k), Cow::Owned(v)))
1745 .collect()
1746 });
1747
1748 Ok(())
1749 }
1750}
1751
1752pub(crate) fn resource_to_pyobject<'p>(
1754 py: Python<'p>,
1755 resource: &Resource<u8>,
1756) -> PyResult<&'p PyCell<OxidizedResource>> {
1757 PyCell::new(
1758 py,
1759 OxidizedResource {
1760 resource: RefCell::new(resource.to_owned()),
1761 },
1762 )
1763}
1764
1765#[inline]
1766pub(crate) fn pyobject_to_resource(resource: &OxidizedResource) -> Resource<'static, u8> {
1767 resource.resource.borrow().clone()
1768}