python_packaging/
resource_collection.rs

1// Copyright 2022 Gregory Szorc.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9/*! Functionality for collecting Python resources. */
10
11use {
12    crate::{
13        bytecode::{
14            compute_bytecode_header, BytecodeHeaderMode, CompileMode, PythonBytecodeCompiler,
15        },
16        libpython::LibPythonBuildContext,
17        licensing::{LicensedComponent, LicensedComponents},
18        location::{AbstractResourceLocation, ConcreteResourceLocation},
19        module_util::{packages_from_module_name, resolve_path_for_module},
20        python_source::has_dunder_file,
21        resource::{
22            BytecodeOptimizationLevel, PythonExtensionModule, PythonModuleBytecode,
23            PythonModuleBytecodeFromSource, PythonModuleSource, PythonPackageDistributionResource,
24            PythonPackageResource, PythonResource, SharedLibrary,
25        },
26    },
27    anyhow::{anyhow, Context, Result},
28    python_packed_resources::Resource,
29    simple_file_manifest::{File, FileData, FileEntry, FileManifest},
30    std::{
31        borrow::Cow,
32        collections::{BTreeMap, BTreeSet, HashMap},
33        path::PathBuf,
34    },
35};
36
37/// Represents a single file install.
38///
39/// Tuple is the relative install path, the data to install, and whether the file
40/// should be executable.
41pub type FileInstall = (PathBuf, FileData, bool);
42
43/// Describes how Python module bytecode will be obtained.
44#[derive(Clone, Debug, PartialEq)]
45pub enum PythonModuleBytecodeProvider {
46    /// Bytecode is already available.
47    Provided(FileData),
48    /// Bytecode will be computed from source.
49    FromSource(FileData),
50}
51
52/// Represents a Python resource entry before it is packaged.
53///
54/// Instances hold the same fields as `Resource` except fields holding
55/// content are backed by a `FileData` instead of `Vec<u8>`, since
56/// we want data resolution to be lazy. In addition, bytecode can either be
57/// provided verbatim or via source.
58#[derive(Clone, Debug, Default, PartialEq)]
59pub struct PrePackagedResource {
60    pub name: String,
61    pub is_package: bool,
62    pub is_namespace_package: bool,
63    pub in_memory_source: Option<FileData>,
64    pub in_memory_bytecode: Option<PythonModuleBytecodeProvider>,
65    pub in_memory_bytecode_opt1: Option<PythonModuleBytecodeProvider>,
66    pub in_memory_bytecode_opt2: Option<PythonModuleBytecodeProvider>,
67    pub in_memory_extension_module_shared_library: Option<FileData>,
68    pub in_memory_resources: Option<BTreeMap<String, FileData>>,
69    pub in_memory_distribution_resources: Option<BTreeMap<String, FileData>>,
70    pub in_memory_shared_library: Option<FileData>,
71    pub shared_library_dependency_names: Option<Vec<String>>,
72    // (prefix, source code)
73    pub relative_path_module_source: Option<(String, FileData)>,
74    // (prefix, bytecode tag, source code)
75    pub relative_path_bytecode: Option<(String, String, PythonModuleBytecodeProvider)>,
76    pub relative_path_bytecode_opt1: Option<(String, String, PythonModuleBytecodeProvider)>,
77    pub relative_path_bytecode_opt2: Option<(String, String, PythonModuleBytecodeProvider)>,
78    // (path, data)
79    pub relative_path_extension_module_shared_library: Option<(PathBuf, FileData)>,
80    pub relative_path_package_resources: Option<BTreeMap<String, (PathBuf, FileData)>>,
81    pub relative_path_distribution_resources: Option<BTreeMap<String, (PathBuf, FileData)>>,
82    pub relative_path_shared_library: Option<(String, PathBuf, FileData)>,
83    pub is_module: bool,
84    pub is_builtin_extension_module: bool,
85    pub is_frozen_module: bool,
86    pub is_extension_module: bool,
87    pub is_shared_library: bool,
88    pub is_utf8_filename_data: bool,
89    pub file_executable: bool,
90    pub file_data_embedded: Option<FileData>,
91    pub file_data_utf8_relative_path: Option<(PathBuf, FileData)>,
92}
93
94impl PrePackagedResource {
95    /// Whether this resource represents a Python resource.
96    pub fn is_python_resource(&self) -> bool {
97        self.is_module
98            || self.is_builtin_extension_module
99            || self.is_frozen_module
100            || self.is_extension_module
101    }
102
103    /// Convert the instance to a `Resource`.
104    ///
105    /// This will compile bytecode from source code using the specified compiler.
106    /// It will also emit a list of file installs that must be performed for all
107    /// referenced resources to function as intended.
108    pub fn to_resource<'a>(
109        &self,
110        compiler: &mut dyn PythonBytecodeCompiler,
111    ) -> Result<(Resource<'a, u8>, Vec<FileInstall>)> {
112        let mut installs = Vec::new();
113
114        let resource = Resource {
115            name: Cow::Owned(self.name.clone()),
116            is_python_package: self.is_package,
117            is_python_namespace_package: self.is_namespace_package,
118            in_memory_source: if let Some(location) = &self.in_memory_source {
119                Some(Cow::Owned(location.resolve_content()?))
120            } else {
121                None
122            },
123            in_memory_bytecode: match &self.in_memory_bytecode {
124                Some(PythonModuleBytecodeProvider::Provided(location)) => {
125                    Some(Cow::Owned(location.resolve_content()?))
126                }
127                Some(PythonModuleBytecodeProvider::FromSource(location)) => Some(Cow::Owned(
128                    compiler
129                        .compile(
130                            &location.resolve_content()?,
131                            &self.name,
132                            BytecodeOptimizationLevel::Zero,
133                            CompileMode::Bytecode,
134                        )
135                        .context("compiling in-memory bytecode")?,
136                )),
137                None => None,
138            },
139            in_memory_bytecode_opt1: match &self.in_memory_bytecode_opt1 {
140                Some(PythonModuleBytecodeProvider::Provided(location)) => {
141                    Some(Cow::Owned(location.resolve_content()?))
142                }
143                Some(PythonModuleBytecodeProvider::FromSource(location)) => Some(Cow::Owned(
144                    compiler
145                        .compile(
146                            &location.resolve_content()?,
147                            &self.name,
148                            BytecodeOptimizationLevel::One,
149                            CompileMode::Bytecode,
150                        )
151                        .context("compiling in-memory bytecode opt-1")?,
152                )),
153                None => None,
154            },
155            in_memory_bytecode_opt2: match &self.in_memory_bytecode_opt2 {
156                Some(PythonModuleBytecodeProvider::Provided(location)) => {
157                    Some(Cow::Owned(location.resolve_content()?))
158                }
159                Some(PythonModuleBytecodeProvider::FromSource(location)) => Some(Cow::Owned(
160                    compiler
161                        .compile(
162                            &location.resolve_content()?,
163                            &self.name,
164                            BytecodeOptimizationLevel::Two,
165                            CompileMode::Bytecode,
166                        )
167                        .context("compiling in-memory bytecode opt2")?,
168                )),
169                None => None,
170            },
171            in_memory_extension_module_shared_library: if let Some(location) =
172                &self.in_memory_extension_module_shared_library
173            {
174                Some(Cow::Owned(location.resolve_content()?))
175            } else {
176                None
177            },
178            in_memory_package_resources: if let Some(resources) = &self.in_memory_resources {
179                let mut res = HashMap::new();
180                for (key, location) in resources {
181                    res.insert(
182                        Cow::Owned(key.clone()),
183                        Cow::Owned(location.resolve_content()?),
184                    );
185                }
186                Some(res)
187            } else {
188                None
189            },
190            in_memory_distribution_resources: if let Some(resources) =
191                &self.in_memory_distribution_resources
192            {
193                let mut res = HashMap::new();
194                for (key, location) in resources {
195                    res.insert(
196                        Cow::Owned(key.clone()),
197                        Cow::Owned(location.resolve_content()?),
198                    );
199                }
200                Some(res)
201            } else {
202                None
203            },
204            in_memory_shared_library: if let Some(location) = &self.in_memory_shared_library {
205                Some(Cow::Owned(location.resolve_content()?))
206            } else {
207                None
208            },
209            shared_library_dependency_names: self
210                .shared_library_dependency_names
211                .as_ref()
212                .map(|x| x.iter().map(|x| Cow::Owned(x.clone())).collect()),
213            relative_path_module_source: if let Some((prefix, location)) =
214                &self.relative_path_module_source
215            {
216                let path = resolve_path_for_module(prefix, &self.name, self.is_package, None);
217
218                installs.push((path.clone(), location.clone(), false));
219
220                Some(Cow::Owned(path))
221            } else {
222                None
223            },
224            relative_path_module_bytecode: if let Some((prefix, cache_tag, provider)) =
225                &self.relative_path_bytecode
226            {
227                let path = resolve_path_for_module(
228                    prefix,
229                    &self.name,
230                    self.is_package,
231                    Some(&format!(
232                        "{}{}",
233                        cache_tag,
234                        BytecodeOptimizationLevel::Zero.to_extra_tag()
235                    )),
236                );
237
238                installs.push((
239                    path.clone(),
240                    FileData::Memory(match provider {
241                        PythonModuleBytecodeProvider::FromSource(location) => compiler
242                            .compile(
243                                &location.resolve_content()?,
244                                &self.name,
245                                BytecodeOptimizationLevel::Zero,
246                                CompileMode::PycUncheckedHash,
247                            )
248                            .context("compiling relative path module bytecode")?,
249                        PythonModuleBytecodeProvider::Provided(location) => {
250                            let mut data = compute_bytecode_header(
251                                compiler.get_magic_number(),
252                                BytecodeHeaderMode::UncheckedHash(0),
253                            )?;
254                            data.extend(location.resolve_content()?);
255
256                            data
257                        }
258                    }),
259                    false,
260                ));
261
262                Some(Cow::Owned(path))
263            } else {
264                None
265            },
266            relative_path_module_bytecode_opt1: if let Some((prefix, cache_tag, provider)) =
267                &self.relative_path_bytecode_opt1
268            {
269                let path = resolve_path_for_module(
270                    prefix,
271                    &self.name,
272                    self.is_package,
273                    Some(&format!(
274                        "{}{}",
275                        cache_tag,
276                        BytecodeOptimizationLevel::One.to_extra_tag()
277                    )),
278                );
279
280                installs.push((
281                    path.clone(),
282                    FileData::Memory(match provider {
283                        PythonModuleBytecodeProvider::FromSource(location) => compiler
284                            .compile(
285                                &location.resolve_content()?,
286                                &self.name,
287                                BytecodeOptimizationLevel::One,
288                                CompileMode::PycUncheckedHash,
289                            )
290                            .context("compiling relative path module bytecode opt-1")?,
291                        PythonModuleBytecodeProvider::Provided(location) => {
292                            let mut data = compute_bytecode_header(
293                                compiler.get_magic_number(),
294                                BytecodeHeaderMode::UncheckedHash(0),
295                            )?;
296                            data.extend(location.resolve_content()?);
297
298                            data
299                        }
300                    }),
301                    false,
302                ));
303
304                Some(Cow::Owned(path))
305            } else {
306                None
307            },
308            relative_path_module_bytecode_opt2: if let Some((prefix, cache_tag, provider)) =
309                &self.relative_path_bytecode_opt2
310            {
311                let path = resolve_path_for_module(
312                    prefix,
313                    &self.name,
314                    self.is_package,
315                    Some(&format!(
316                        "{}{}",
317                        cache_tag,
318                        BytecodeOptimizationLevel::Two.to_extra_tag()
319                    )),
320                );
321
322                installs.push((
323                    path.clone(),
324                    FileData::Memory(match provider {
325                        PythonModuleBytecodeProvider::FromSource(location) => compiler.compile(
326                            &location.resolve_content()?,
327                            &self.name,
328                            BytecodeOptimizationLevel::Two,
329                            CompileMode::PycUncheckedHash,
330                        )?,
331                        PythonModuleBytecodeProvider::Provided(location) => {
332                            let mut data = compute_bytecode_header(
333                                compiler.get_magic_number(),
334                                BytecodeHeaderMode::UncheckedHash(0),
335                            )
336                            .context("compiling relative path module bytecode opt-2")?;
337                            data.extend(location.resolve_content()?);
338
339                            data
340                        }
341                    }),
342                    false,
343                ));
344
345                Some(Cow::Owned(path))
346            } else {
347                None
348            },
349            relative_path_extension_module_shared_library: if let Some((path, location)) =
350                &self.relative_path_extension_module_shared_library
351            {
352                installs.push((path.clone(), location.clone(), true));
353
354                Some(Cow::Owned(path.clone()))
355            } else {
356                None
357            },
358            relative_path_package_resources: if let Some(resources) =
359                &self.relative_path_package_resources
360            {
361                let mut res = HashMap::new();
362                for (key, (path, location)) in resources {
363                    installs.push((path.clone(), location.clone(), false));
364
365                    res.insert(Cow::Owned(key.clone()), Cow::Owned(path.clone()));
366                }
367                Some(res)
368            } else {
369                None
370            },
371            relative_path_distribution_resources: if let Some(resources) =
372                &self.relative_path_distribution_resources
373            {
374                let mut res = HashMap::new();
375                for (key, (path, location)) in resources {
376                    installs.push((path.clone(), location.clone(), false));
377
378                    res.insert(Cow::Owned(key.clone()), Cow::Owned(path.clone()));
379                }
380                Some(res)
381            } else {
382                None
383            },
384            is_python_module: self.is_module,
385            is_python_builtin_extension_module: self.is_builtin_extension_module,
386            is_python_frozen_module: self.is_frozen_module,
387            is_python_extension_module: self.is_extension_module,
388            is_shared_library: self.is_shared_library,
389            is_utf8_filename_data: self.is_utf8_filename_data,
390            file_executable: self.file_executable,
391            file_data_embedded: if let Some(location) = &self.file_data_embedded {
392                Some(Cow::Owned(location.resolve_content()?))
393            } else {
394                None
395            },
396            file_data_utf8_relative_path: if let Some((path, location)) =
397                &self.file_data_utf8_relative_path
398            {
399                installs.push((path.clone(), location.clone(), self.file_executable));
400
401                Some(Cow::Owned(path.to_string_lossy().to_string()))
402            } else {
403                None
404            },
405        };
406
407        if let Some((prefix, filename, location)) = &self.relative_path_shared_library {
408            installs.push((PathBuf::from(prefix).join(filename), location.clone(), true));
409        }
410
411        Ok((resource, installs))
412    }
413}
414
415/// Fill in missing data on parent packages.
416///
417/// When resources are added, their parent packages could be missing
418/// data. If we simply materialized the child resources without the
419/// parents, Python's importer would get confused due to the missing
420/// resources.
421///
422/// This function fills in the blanks in our resources state.
423///
424/// The way this works is that if a child resource has data in
425/// a particular field, we populate that field in all its parent
426/// packages. If a corresponding fields is already populated, we
427/// copy its data as well.
428pub fn populate_parent_packages(
429    resources: &mut BTreeMap<String, PrePackagedResource>,
430) -> Result<()> {
431    let original_resources = resources
432        .iter()
433        .filter_map(|(k, v)| {
434            if v.is_python_resource() {
435                Some((k.to_owned(), v.to_owned()))
436            } else {
437                None
438            }
439        })
440        .collect::<Vec<(String, PrePackagedResource)>>();
441
442    for (name, original) in original_resources {
443        for package in packages_from_module_name(&name) {
444            let entry = resources
445                .entry(package.clone())
446                .or_insert_with(|| PrePackagedResource {
447                    name: package,
448                    ..PrePackagedResource::default()
449                });
450
451            // Parents must be modules + packages by definition.
452            entry.is_module = true;
453            entry.is_package = true;
454
455            // We want to materialize bytecode on parent packages no matter
456            // what. If the original resource has a variant of bytecode in a
457            // location, we materialize that variant on parents. We take
458            // the source from the parent resource, if present. Otherwise
459            // defaulting to empty.
460            if original.in_memory_bytecode.is_some() && entry.in_memory_bytecode.is_none() {
461                entry.in_memory_bytecode = Some(PythonModuleBytecodeProvider::FromSource(
462                    if let Some(source) = &entry.in_memory_source {
463                        source.clone()
464                    } else {
465                        FileData::Memory(vec![])
466                    },
467                ));
468            }
469            if original.in_memory_bytecode_opt1.is_some() && entry.in_memory_bytecode_opt1.is_none()
470            {
471                entry.in_memory_bytecode_opt1 = Some(PythonModuleBytecodeProvider::FromSource(
472                    if let Some(source) = &entry.in_memory_source {
473                        source.clone()
474                    } else {
475                        FileData::Memory(vec![])
476                    },
477                ));
478            }
479            if original.in_memory_bytecode_opt2.is_some() && entry.in_memory_bytecode_opt2.is_none()
480            {
481                entry.in_memory_bytecode_opt2 = Some(PythonModuleBytecodeProvider::FromSource(
482                    if let Some(source) = &entry.in_memory_source {
483                        source.clone()
484                    } else {
485                        FileData::Memory(vec![])
486                    },
487                ));
488            }
489
490            if let Some((prefix, cache_tag, _)) = &original.relative_path_bytecode {
491                if entry.relative_path_bytecode.is_none() {
492                    entry.relative_path_bytecode = Some((
493                        prefix.clone(),
494                        cache_tag.clone(),
495                        PythonModuleBytecodeProvider::FromSource(
496                            if let Some((_, location)) = &entry.relative_path_module_source {
497                                location.clone()
498                            } else {
499                                FileData::Memory(vec![])
500                            },
501                        ),
502                    ));
503                }
504            }
505
506            if let Some((prefix, cache_tag, _)) = &original.relative_path_bytecode_opt1 {
507                if entry.relative_path_bytecode_opt1.is_none() {
508                    entry.relative_path_bytecode_opt1 = Some((
509                        prefix.clone(),
510                        cache_tag.clone(),
511                        PythonModuleBytecodeProvider::FromSource(
512                            if let Some((_, location)) = &entry.relative_path_module_source {
513                                location.clone()
514                            } else {
515                                FileData::Memory(vec![])
516                            },
517                        ),
518                    ));
519                }
520            }
521
522            if let Some((prefix, cache_tag, _)) = &original.relative_path_bytecode_opt2 {
523                if entry.relative_path_bytecode_opt2.is_none() {
524                    entry.relative_path_bytecode_opt2 = Some((
525                        prefix.clone(),
526                        cache_tag.clone(),
527                        PythonModuleBytecodeProvider::FromSource(
528                            if let Some((_, location)) = &entry.relative_path_module_source {
529                                location.clone()
530                            } else {
531                                FileData::Memory(vec![])
532                            },
533                        ),
534                    ));
535                }
536            }
537
538            // If the child had path-based source, we need to materialize source as well.
539            if let Some((prefix, _)) = &original.relative_path_module_source {
540                entry
541                    .relative_path_module_source
542                    .get_or_insert_with(|| (prefix.clone(), FileData::Memory(vec![])));
543            }
544
545            // Ditto for in-memory source.
546            if original.in_memory_source.is_some() {
547                entry
548                    .in_memory_source
549                    .get_or_insert(FileData::Memory(vec![]));
550            }
551        }
552    }
553
554    Ok(())
555}
556
557/// Defines how a Python resource should be added to a `PythonResourceCollector`.
558#[derive(Clone, Debug, PartialEq, Eq)]
559pub struct PythonResourceAddCollectionContext {
560    /// Whether the resource should be included in `PythonResourceCollection`.
561    pub include: bool,
562
563    /// The location the resource should be loaded from.
564    pub location: ConcreteResourceLocation,
565
566    /// Optional fallback location from which to load the resource from.
567    ///
568    /// If adding the resource to `location` fails, and this is defined,
569    /// we will fall back to adding the resource to this location.
570    pub location_fallback: Option<ConcreteResourceLocation>,
571
572    /// Whether to store Python source code for a `PythonModuleSource`.
573    ///
574    /// When handling a `PythonModuleSource`, sometimes you want to
575    /// write just bytecode or source + bytecode. This flags allows
576    /// controlling this behavior.
577    pub store_source: bool,
578
579    /// Whether to store Python bytecode for optimization level 0.
580    pub optimize_level_zero: bool,
581
582    /// Whether to store Python bytecode for optimization level 1.
583    pub optimize_level_one: bool,
584
585    /// Whether to store Python bytecode for optimization level 2.
586    pub optimize_level_two: bool,
587}
588
589impl PythonResourceAddCollectionContext {
590    /// Replace the content of `self` with content of `other`.
591    pub fn replace(&mut self, other: &Self) {
592        self.include = other.include;
593        self.location = other.location.clone();
594        self.location_fallback = other.location_fallback.clone();
595        self.store_source = other.store_source;
596        self.optimize_level_zero = other.optimize_level_zero;
597        self.optimize_level_one = other.optimize_level_one;
598        self.optimize_level_two = other.optimize_level_two;
599    }
600}
601
602/// Describes the result of adding a resource to a collector.
603#[derive(Clone, Debug)]
604pub enum AddResourceAction {
605    /// Resource with specified description wasn't added because add_include=false.
606    NoInclude(String),
607    /// Resource not added because of Python bytecode optimization level mismatch.
608    BytecodeOptimizationLevelMismatch(String),
609    /// Resource with specified description was added to the specified location.
610    Added(String, ConcreteResourceLocation),
611    /// Built-in Python extension module.
612    AddedBuiltinExtensionModule(String),
613}
614
615impl ToString for AddResourceAction {
616    fn to_string(&self) -> String {
617        match self {
618            Self::NoInclude(name) => {
619                format!("ignored adding {} because fails inclusion filter", name)
620            }
621            Self::BytecodeOptimizationLevelMismatch(name) => {
622                format!("ignored adding Python module bytecode for {} because of optimization level mismatch", name)
623            }
624            Self::Added(name, location) => {
625                format!("added {} to {}", name, location.to_string())
626            }
627            Self::AddedBuiltinExtensionModule(name) => {
628                format!("added builtin Python extension module {}", name)
629            }
630        }
631    }
632}
633
634/// Represents a finalized collection of Python resources.
635///
636/// Instances are produced from a `PythonResourceCollector` and a
637/// `PythonBytecodeCompiler` to produce bytecode.
638#[derive(Clone, Debug, Default)]
639pub struct CompiledResourcesCollection<'a> {
640    /// All indexes resources.
641    pub resources: BTreeMap<String, Resource<'a, u8>>,
642
643    /// Extra file installs that must be performed so referenced files are available.
644    pub extra_files: Vec<FileInstall>,
645}
646
647impl<'a> CompiledResourcesCollection<'a> {
648    /// Write resources to packed resources data, version 1.
649    pub fn write_packed_resources<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
650        python_packed_resources::write_packed_resources_v3(
651            &self
652                .resources
653                .values()
654                .cloned()
655                .collect::<Vec<Resource<'a, u8>>>(),
656            writer,
657            None,
658        )
659    }
660
661    /// Convert the file installs to a [FileManifest].
662    pub fn extra_files_manifest(&self) -> Result<FileManifest> {
663        let mut m = FileManifest::default();
664
665        for (path, location, executable) in &self.extra_files {
666            m.add_file_entry(
667                path,
668                FileEntry::new_from_data(location.resolve_content()?, *executable),
669            )?;
670        }
671
672        Ok(m)
673    }
674}
675
676/// Type used to collect Python resources so they can be serialized.
677///
678/// We often want to turn Python resource primitives (module source,
679/// bytecode, etc) into a collection of `Resource` so they can be
680/// serialized to the *Python packed resources* format. This type
681/// exists to facilitate doing this.
682///
683/// This type is not only responsible for tracking resources but also for
684/// enforcing policies on where those resources can be loaded from and
685/// what types of resources are allowed. This includes tracking the
686/// licensing metadata for indexed resources.
687#[derive(Debug, Clone)]
688pub struct PythonResourceCollector {
689    /// Where resources can be placed.
690    allowed_locations: Vec<AbstractResourceLocation>,
691
692    /// Allowed locations for extension modules.
693    ///
694    /// This is applied in addition to `allowed_locations` and can be
695    /// more strict.
696    allowed_extension_module_locations: Vec<AbstractResourceLocation>,
697
698    /// Whether builtin extension modules outside the standard library are allowed.
699    ///
700    /// This is effectively "are we building a custom libpython." If true,
701    /// we can take object files / static libraries from adding extension
702    /// modules are add the extension module as a built-in. If false, only
703    /// builtin extension modules already in libpython can be added as a
704    /// built-in.
705    allow_new_builtin_extension_modules: bool,
706
707    /// Whether untyped files (`File`) can be added.
708    allow_files: bool,
709
710    /// Named resources that have been collected.
711    resources: BTreeMap<String, PrePackagedResource>,
712
713    /// Collection of software components which are licensed.
714    licensed_components: LicensedComponents,
715}
716
717impl PythonResourceCollector {
718    /// Construct a new instance of the collector.
719    ///
720    /// The instance is associated with a resources policy to validate that
721    /// added resources conform with rules.
722    ///
723    /// We also pass a Python bytecode cache tag, which is used to
724    /// derive filenames.
725    pub fn new(
726        allowed_locations: Vec<AbstractResourceLocation>,
727        allowed_extension_module_locations: Vec<AbstractResourceLocation>,
728        allow_new_builtin_extension_modules: bool,
729        allow_files: bool,
730    ) -> Self {
731        Self {
732            allowed_locations,
733            allowed_extension_module_locations,
734            allow_new_builtin_extension_modules,
735            allow_files,
736            resources: BTreeMap::new(),
737            licensed_components: LicensedComponents::default(),
738        }
739    }
740
741    /// Obtain locations that resources can be loaded from.
742    pub fn allowed_locations(&self) -> &Vec<AbstractResourceLocation> {
743        &self.allowed_locations
744    }
745
746    /// Obtain a set of all top-level Python module names registered with the collector.
747    ///
748    /// The returned values correspond to packages or single file modules without
749    /// children modules.
750    pub fn all_top_level_module_names(&self) -> BTreeSet<String> {
751        self.resources
752            .values()
753            .filter_map(|r| {
754                if r.is_python_resource() {
755                    let name = if let Some(idx) = r.name.find('.') {
756                        &r.name[0..idx]
757                    } else {
758                        &r.name
759                    };
760
761                    Some(name.to_string())
762                } else {
763                    None
764                }
765            })
766            .collect::<BTreeSet<_>>()
767    }
768
769    /// Validate that a resource add in the specified location is allowed.
770    pub fn check_policy(&self, location: AbstractResourceLocation) -> Result<()> {
771        if self.allowed_locations.contains(&location) {
772            Ok(())
773        } else {
774            Err(anyhow!(
775                "resource collector does not allow resources in {}",
776                (&location).to_string()
777            ))
778        }
779    }
780
781    /// Apply a filter function on resources in this collection and mutate in place.
782    ///
783    /// If the filter function returns true, the item will be preserved.
784    pub fn filter_resources_mut<F>(&mut self, filter: F) -> Result<()>
785    where
786        F: Fn(&PrePackagedResource) -> bool,
787    {
788        self.resources = self
789            .resources
790            .iter()
791            .filter_map(|(k, v)| {
792                if filter(v) {
793                    Some((k.clone(), v.clone()))
794                } else {
795                    None
796                }
797            })
798            .collect();
799
800        Ok(())
801    }
802
803    /// Obtain an iterator over the resources in this collector.
804    pub fn iter_resources(&self) -> impl Iterator<Item = (&String, &PrePackagedResource)> {
805        Box::new(self.resources.iter())
806    }
807
808    /// Register a licensed software component to this collection.
809    pub fn add_licensed_component(&mut self, component: LicensedComponent) -> Result<()> {
810        self.licensed_components.add_component(component);
811
812        Ok(())
813    }
814
815    /// Obtain a finalized collection of licensed components.
816    ///
817    /// The collection has entries for components that lack licenses and has additional
818    /// normalization performed.
819    pub fn normalized_licensed_components(&self) -> LicensedComponents {
820        self.licensed_components.normalize_python_modules()
821    }
822
823    /// Add Python module source with a specific location.
824    pub fn add_python_module_source(
825        &mut self,
826        module: &PythonModuleSource,
827        location: &ConcreteResourceLocation,
828    ) -> Result<Vec<AddResourceAction>> {
829        self.check_policy(location.into())?;
830
831        let entry = self
832            .resources
833            .entry(module.name.clone())
834            .or_insert_with(|| PrePackagedResource {
835                name: module.name.clone(),
836                ..PrePackagedResource::default()
837            });
838
839        entry.is_module = true;
840        entry.is_package = module.is_package;
841
842        match location {
843            ConcreteResourceLocation::InMemory => {
844                entry.in_memory_source = Some(module.source.clone());
845            }
846            ConcreteResourceLocation::RelativePath(prefix) => {
847                entry.relative_path_module_source =
848                    Some((prefix.to_string(), module.source.clone()));
849            }
850        }
851
852        Ok(vec![AddResourceAction::Added(
853            module.description(),
854            location.clone(),
855        )])
856    }
857
858    /// Add Python module source using an add context to influence operation.
859    ///
860    /// All of the context's properties are respected. This includes doing
861    /// nothing if `include` is false, not adding source if `store_source` is
862    /// false, and automatically deriving a bytecode request if the
863    /// `optimize_level_*` fields are set.
864    ///
865    /// This method is a glorified proxy to other `add_*` methods: it
866    /// simply contains the logic for expanding the context's wishes into
867    /// function calls.
868    pub fn add_python_module_source_with_context(
869        &mut self,
870        module: &PythonModuleSource,
871        add_context: &PythonResourceAddCollectionContext,
872    ) -> Result<Vec<AddResourceAction>> {
873        if !add_context.include {
874            return Ok(vec![AddResourceAction::NoInclude(module.description())]);
875        }
876
877        let mut actions = vec![];
878
879        if add_context.store_source {
880            actions.extend(self.add_python_resource_with_locations(
881                &module.into(),
882                &add_context.location,
883                &add_context.location_fallback,
884            )?);
885        }
886
887        // Derive bytecode as requested.
888        if add_context.optimize_level_zero {
889            actions.extend(
890                self.add_python_resource_with_locations(
891                    &module
892                        .as_bytecode_module(BytecodeOptimizationLevel::Zero)
893                        .into(),
894                    &add_context.location,
895                    &add_context.location_fallback,
896                )?,
897            );
898        }
899
900        if add_context.optimize_level_one {
901            actions.extend(
902                self.add_python_resource_with_locations(
903                    &module
904                        .as_bytecode_module(BytecodeOptimizationLevel::One)
905                        .into(),
906                    &add_context.location,
907                    &add_context.location_fallback,
908                )?,
909            );
910        }
911
912        if add_context.optimize_level_two {
913            actions.extend(
914                self.add_python_resource_with_locations(
915                    &module
916                        .as_bytecode_module(BytecodeOptimizationLevel::Two)
917                        .into(),
918                    &add_context.location,
919                    &add_context.location_fallback,
920                )?,
921            );
922        }
923
924        Ok(actions)
925    }
926
927    /// Add Python module bytecode to the specified location.
928    pub fn add_python_module_bytecode(
929        &mut self,
930        module: &PythonModuleBytecode,
931        location: &ConcreteResourceLocation,
932    ) -> Result<Vec<AddResourceAction>> {
933        self.check_policy(location.into())?;
934
935        let entry = self
936            .resources
937            .entry(module.name.clone())
938            .or_insert_with(|| PrePackagedResource {
939                name: module.name.clone(),
940                ..PrePackagedResource::default()
941            });
942
943        entry.is_module = true;
944        entry.is_package = module.is_package;
945
946        // TODO having to resolve the FileData here is a bit unfortunate.
947        // We could invent a better type to allow the I/O to remain lazy.
948        let bytecode =
949            PythonModuleBytecodeProvider::Provided(FileData::Memory(module.resolve_bytecode()?));
950
951        match location {
952            ConcreteResourceLocation::InMemory => match module.optimize_level {
953                BytecodeOptimizationLevel::Zero => {
954                    entry.in_memory_bytecode = Some(bytecode);
955                }
956                BytecodeOptimizationLevel::One => {
957                    entry.in_memory_bytecode_opt1 = Some(bytecode);
958                }
959                BytecodeOptimizationLevel::Two => {
960                    entry.in_memory_bytecode_opt2 = Some(bytecode);
961                }
962            },
963            ConcreteResourceLocation::RelativePath(prefix) => match module.optimize_level {
964                BytecodeOptimizationLevel::Zero => {
965                    entry.relative_path_bytecode =
966                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode));
967                }
968                BytecodeOptimizationLevel::One => {
969                    entry.relative_path_bytecode_opt1 =
970                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode));
971                }
972                BytecodeOptimizationLevel::Two => {
973                    entry.relative_path_bytecode_opt2 =
974                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode));
975                }
976            },
977        }
978
979        Ok(vec![AddResourceAction::Added(
980            module.description(),
981            location.clone(),
982        )])
983    }
984
985    /// Add Python module bytecode using an add context.
986    ///
987    /// This takes the context's fields into consideration when adding
988    /// the resource. If `include` is false, this is a no-op. The context
989    /// must also have an `optimize_level_*` field set corresponding with
990    /// the optimization level of the passed bytecode, or this is a no-op.
991    pub fn add_python_module_bytecode_with_context(
992        &mut self,
993        module: &PythonModuleBytecode,
994        add_context: &PythonResourceAddCollectionContext,
995    ) -> Result<Vec<AddResourceAction>> {
996        if !add_context.include {
997            return Ok(vec![AddResourceAction::NoInclude(module.description())]);
998        }
999
1000        match module.optimize_level {
1001            BytecodeOptimizationLevel::Zero => {
1002                if add_context.optimize_level_zero {
1003                    self.add_python_resource_with_locations(
1004                        &module.into(),
1005                        &add_context.location,
1006                        &add_context.location_fallback,
1007                    )
1008                } else {
1009                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1010                        module.name.clone(),
1011                    )])
1012                }
1013            }
1014            BytecodeOptimizationLevel::One => {
1015                if add_context.optimize_level_one {
1016                    self.add_python_resource_with_locations(
1017                        &module.into(),
1018                        &add_context.location,
1019                        &add_context.location_fallback,
1020                    )
1021                } else {
1022                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1023                        module.name.clone(),
1024                    )])
1025                }
1026            }
1027            BytecodeOptimizationLevel::Two => {
1028                if add_context.optimize_level_two {
1029                    self.add_python_resource_with_locations(
1030                        &module.into(),
1031                        &add_context.location,
1032                        &add_context.location_fallback,
1033                    )
1034                } else {
1035                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1036                        module.name.clone(),
1037                    )])
1038                }
1039            }
1040        }
1041    }
1042
1043    /// Add Python module bytecode derived from source code to the collection.
1044    pub fn add_python_module_bytecode_from_source(
1045        &mut self,
1046        module: &PythonModuleBytecodeFromSource,
1047        location: &ConcreteResourceLocation,
1048    ) -> Result<Vec<AddResourceAction>> {
1049        self.check_policy(location.into())?;
1050
1051        let entry = self
1052            .resources
1053            .entry(module.name.clone())
1054            .or_insert_with(|| PrePackagedResource {
1055                name: module.name.clone(),
1056                ..PrePackagedResource::default()
1057            });
1058
1059        entry.is_module = true;
1060        entry.is_package = module.is_package;
1061
1062        let bytecode = PythonModuleBytecodeProvider::FromSource(module.source.clone());
1063
1064        match location {
1065            ConcreteResourceLocation::InMemory => match module.optimize_level {
1066                BytecodeOptimizationLevel::Zero => {
1067                    entry.in_memory_bytecode = Some(bytecode);
1068                }
1069                BytecodeOptimizationLevel::One => {
1070                    entry.in_memory_bytecode_opt1 = Some(bytecode);
1071                }
1072                BytecodeOptimizationLevel::Two => {
1073                    entry.in_memory_bytecode_opt2 = Some(bytecode);
1074                }
1075            },
1076            ConcreteResourceLocation::RelativePath(prefix) => match module.optimize_level {
1077                BytecodeOptimizationLevel::Zero => {
1078                    entry.relative_path_bytecode =
1079                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode))
1080                }
1081                BytecodeOptimizationLevel::One => {
1082                    entry.relative_path_bytecode_opt1 =
1083                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode))
1084                }
1085                BytecodeOptimizationLevel::Two => {
1086                    entry.relative_path_bytecode_opt2 =
1087                        Some((prefix.to_string(), module.cache_tag.clone(), bytecode))
1088                }
1089            },
1090        }
1091
1092        Ok(vec![AddResourceAction::Added(
1093            module.description(),
1094            location.clone(),
1095        )])
1096    }
1097
1098    /// Add Python module bytecode from source using an add context to influence operations.
1099    ///
1100    /// This method respects the settings of the context, including `include`
1101    /// and the `optimize_level_*` fields.
1102    ///
1103    /// `PythonModuleBytecodeFromSource` defines an explicit bytecode optimization level,
1104    /// so this method can result in at most 1 bytecode request being added to the
1105    /// collection.
1106    pub fn add_python_module_bytecode_from_source_with_context(
1107        &mut self,
1108        module: &PythonModuleBytecodeFromSource,
1109        add_context: &PythonResourceAddCollectionContext,
1110    ) -> Result<Vec<AddResourceAction>> {
1111        if !add_context.include {
1112            return Ok(vec![AddResourceAction::NoInclude(module.description())]);
1113        }
1114
1115        match module.optimize_level {
1116            BytecodeOptimizationLevel::Zero => {
1117                if add_context.optimize_level_zero {
1118                    self.add_python_resource_with_locations(
1119                        &module.into(),
1120                        &add_context.location,
1121                        &add_context.location_fallback,
1122                    )
1123                } else {
1124                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1125                        module.name.clone(),
1126                    )])
1127                }
1128            }
1129            BytecodeOptimizationLevel::One => {
1130                if add_context.optimize_level_one {
1131                    self.add_python_resource_with_locations(
1132                        &module.into(),
1133                        &add_context.location,
1134                        &add_context.location_fallback,
1135                    )
1136                } else {
1137                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1138                        module.name.clone(),
1139                    )])
1140                }
1141            }
1142            BytecodeOptimizationLevel::Two => {
1143                if add_context.optimize_level_two {
1144                    self.add_python_resource_with_locations(
1145                        &module.into(),
1146                        &add_context.location,
1147                        &add_context.location_fallback,
1148                    )
1149                } else {
1150                    Ok(vec![AddResourceAction::BytecodeOptimizationLevelMismatch(
1151                        module.name.clone(),
1152                    )])
1153                }
1154            }
1155        }
1156    }
1157
1158    /// Add resource data to a given location.
1159    ///
1160    /// Resource data belongs to a Python package and has a name and bytes data.
1161    pub fn add_python_package_resource(
1162        &mut self,
1163        resource: &PythonPackageResource,
1164        location: &ConcreteResourceLocation,
1165    ) -> Result<Vec<AddResourceAction>> {
1166        self.check_policy(location.into())?;
1167
1168        let entry = self
1169            .resources
1170            .entry(resource.leaf_package.clone())
1171            .or_insert_with(|| PrePackagedResource {
1172                name: resource.leaf_package.clone(),
1173                ..PrePackagedResource::default()
1174            });
1175
1176        // Adding a resource automatically makes the module a package.
1177        entry.is_module = true;
1178        entry.is_package = true;
1179
1180        match location {
1181            ConcreteResourceLocation::InMemory => {
1182                if entry.in_memory_resources.is_none() {
1183                    entry.in_memory_resources = Some(BTreeMap::new());
1184                }
1185                entry
1186                    .in_memory_resources
1187                    .as_mut()
1188                    .unwrap()
1189                    .insert(resource.relative_name.clone(), resource.data.clone());
1190            }
1191            ConcreteResourceLocation::RelativePath(prefix) => {
1192                if entry.relative_path_package_resources.is_none() {
1193                    entry.relative_path_package_resources = Some(BTreeMap::new());
1194                }
1195
1196                entry
1197                    .relative_path_package_resources
1198                    .as_mut()
1199                    .unwrap()
1200                    .insert(
1201                        resource.relative_name.clone(),
1202                        (resource.resolve_path(prefix), resource.data.clone()),
1203                    );
1204            }
1205        }
1206
1207        Ok(vec![AddResourceAction::Added(
1208            resource.description(),
1209            location.clone(),
1210        )])
1211    }
1212
1213    /// Add a Python package resource using an add context.
1214    ///
1215    /// The fields from the context will be respected. This includes not doing
1216    /// anything if `include` is false.
1217    pub fn add_python_package_resource_with_context(
1218        &mut self,
1219        resource: &PythonPackageResource,
1220        add_context: &PythonResourceAddCollectionContext,
1221    ) -> Result<Vec<AddResourceAction>> {
1222        if !add_context.include {
1223            return Ok(vec![AddResourceAction::NoInclude(resource.description())]);
1224        }
1225
1226        self.add_python_resource_with_locations(
1227            &resource.into(),
1228            &add_context.location,
1229            &add_context.location_fallback,
1230        )
1231    }
1232
1233    /// Add a Python package distribution resource to a given location.
1234    pub fn add_python_package_distribution_resource(
1235        &mut self,
1236        resource: &PythonPackageDistributionResource,
1237        location: &ConcreteResourceLocation,
1238    ) -> Result<Vec<AddResourceAction>> {
1239        self.check_policy(location.into())?;
1240
1241        let entry = self
1242            .resources
1243            .entry(resource.package.clone())
1244            .or_insert_with(|| PrePackagedResource {
1245                name: resource.package.clone(),
1246                ..PrePackagedResource::default()
1247            });
1248
1249        // A distribution resource makes the entity a package.
1250        entry.is_module = true;
1251        entry.is_package = true;
1252
1253        match location {
1254            ConcreteResourceLocation::InMemory => {
1255                if entry.in_memory_distribution_resources.is_none() {
1256                    entry.in_memory_distribution_resources = Some(BTreeMap::new());
1257                }
1258
1259                entry
1260                    .in_memory_distribution_resources
1261                    .as_mut()
1262                    .unwrap()
1263                    .insert(resource.name.clone(), resource.data.clone());
1264            }
1265            ConcreteResourceLocation::RelativePath(prefix) => {
1266                if entry.relative_path_distribution_resources.is_none() {
1267                    entry.relative_path_distribution_resources = Some(BTreeMap::new());
1268                }
1269
1270                entry
1271                    .relative_path_distribution_resources
1272                    .as_mut()
1273                    .unwrap()
1274                    .insert(
1275                        resource.name.clone(),
1276                        (resource.resolve_path(prefix), resource.data.clone()),
1277                    );
1278            }
1279        }
1280
1281        Ok(vec![AddResourceAction::Added(
1282            resource.description(),
1283            location.clone(),
1284        )])
1285    }
1286
1287    /// Add a Python package distribution resource using an add context.
1288    ///
1289    /// The fields from the context will be respected. This includes not doing
1290    /// anything if `include` is false.
1291    pub fn add_python_package_distribution_resource_with_context(
1292        &mut self,
1293        resource: &PythonPackageDistributionResource,
1294        add_context: &PythonResourceAddCollectionContext,
1295    ) -> Result<Vec<AddResourceAction>> {
1296        if !add_context.include {
1297            return Ok(vec![AddResourceAction::NoInclude(resource.description())]);
1298        }
1299
1300        self.add_python_resource_with_locations(
1301            &resource.into(),
1302            &add_context.location,
1303            &add_context.location_fallback,
1304        )
1305    }
1306
1307    /// Add a Python extension module using an add context.
1308    #[allow(clippy::if_same_then_else)]
1309    pub fn add_python_extension_module_with_context(
1310        &mut self,
1311        extension_module: &PythonExtensionModule,
1312        add_context: &PythonResourceAddCollectionContext,
1313    ) -> Result<(Vec<AddResourceAction>, Option<LibPythonBuildContext>)> {
1314        // TODO consult this attribute (it isn't set for built-ins for some reason)
1315        //if !add_context.include {
1316        //    return Ok(None);
1317        // }
1318
1319        // Whether we can load extension modules as standalone shared library files.
1320        let can_load_standalone = self
1321            .allowed_extension_module_locations
1322            .contains(&AbstractResourceLocation::RelativePath);
1323
1324        // Whether we can load extension module dynamic libraries from memory. This
1325        // means we have a dynamic library extension module and that library is loaded
1326        // from memory: this is not a built-in extension!
1327        let can_load_dynamic_library_memory = self
1328            .allowed_extension_module_locations
1329            .contains(&AbstractResourceLocation::InMemory);
1330
1331        // Whether we can link the extension as a built-in. This requires one of the
1332        // following:
1333        //
1334        // 1. The extension module is already a built-in.
1335        // 2. Object files or static library data and for the policy to allow new built-in
1336        //    extension modules.
1337        let can_link_builtin = if extension_module.in_libpython() {
1338            true
1339        } else {
1340            self.allow_new_builtin_extension_modules
1341                && !extension_module.object_file_data.is_empty()
1342        };
1343
1344        // Whether we can produce a standalone shared library extension module.
1345        // TODO consider allowing this if object files are present.
1346        let can_link_standalone = extension_module.shared_library.is_some();
1347
1348        let mut relative_path = if let Some(location) = &add_context.location_fallback {
1349            match location {
1350                ConcreteResourceLocation::RelativePath(ref prefix) => Some(prefix.clone()),
1351                ConcreteResourceLocation::InMemory => None,
1352            }
1353        } else {
1354            None
1355        };
1356
1357        let prefer_in_memory = add_context.location == ConcreteResourceLocation::InMemory;
1358        let prefer_filesystem = match &add_context.location {
1359            ConcreteResourceLocation::RelativePath(_) => true,
1360            ConcreteResourceLocation::InMemory => false,
1361        };
1362
1363        let fallback_in_memory =
1364            add_context.location_fallback == Some(ConcreteResourceLocation::InMemory);
1365        let fallback_filesystem = matches!(
1366            &add_context.location_fallback,
1367            Some(ConcreteResourceLocation::RelativePath(_))
1368        );
1369
1370        // TODO support this.
1371        if prefer_filesystem && fallback_in_memory {
1372            return Err(anyhow!("a preferred location of the filesystem and a fallback from memory is not supported"));
1373        }
1374
1375        let require_in_memory =
1376            prefer_in_memory && (add_context.location_fallback.is_none() || fallback_in_memory);
1377        let require_filesystem =
1378            prefer_filesystem && (add_context.location_fallback.is_none() || fallback_filesystem);
1379
1380        match &add_context.location {
1381            ConcreteResourceLocation::RelativePath(prefix) => {
1382                relative_path = Some(prefix.clone());
1383            }
1384            ConcreteResourceLocation::InMemory => {}
1385        }
1386
1387        // We produce a builtin extension module (by linking object files) if any
1388        // of the following conditions are met:
1389        //
1390        // We are a stdlib extension module built into libpython core
1391        let produce_builtin = if extension_module.is_stdlib && extension_module.builtin_default {
1392            true
1393        // Builtin linking is the only mechanism available to us.
1394        } else if can_link_builtin && (!can_link_standalone || !can_load_standalone) {
1395            true
1396        // We want in memory loading and we can link a builtin
1397        } else {
1398            prefer_in_memory && can_link_builtin && !require_filesystem
1399        };
1400
1401        if require_in_memory && (!can_link_builtin && !can_load_dynamic_library_memory) {
1402            return Err(anyhow!(
1403                "extension module {} cannot be loaded from memory but memory loading required",
1404                extension_module.name
1405            ));
1406        }
1407
1408        if require_filesystem && !can_link_standalone && !produce_builtin {
1409            return Err(anyhow!("extension module {} cannot be materialized as a shared library extension but filesystem loading required", extension_module.name));
1410        }
1411
1412        if !produce_builtin && !can_load_standalone {
1413            return Err(anyhow!("extension module {} cannot be materialized as a shared library because distribution does not support loading extension module shared libraries", extension_module.name));
1414        }
1415
1416        if produce_builtin {
1417            let mut build_context = LibPythonBuildContext::default();
1418
1419            for depends in &extension_module.link_libraries {
1420                if depends.framework {
1421                    build_context.frameworks.insert(depends.name.clone());
1422                } else if depends.system {
1423                    build_context.system_libraries.insert(depends.name.clone());
1424                } else if depends.static_library.is_some() {
1425                    build_context.static_libraries.insert(depends.name.clone());
1426                } else if depends.dynamic_library.is_some() {
1427                    build_context.dynamic_libraries.insert(depends.name.clone());
1428                }
1429            }
1430
1431            if let Some(component) = &extension_module.license {
1432                build_context
1433                    .licensed_components
1434                    .add_component(component.clone());
1435            }
1436
1437            if let Some(init_fn) = &extension_module.init_fn {
1438                build_context
1439                    .init_functions
1440                    .insert(extension_module.name.clone(), init_fn.clone());
1441            }
1442
1443            for location in &extension_module.object_file_data {
1444                build_context.object_files.push(location.clone());
1445            }
1446
1447            let actions = self.add_builtin_python_extension_module(extension_module)?;
1448
1449            Ok((actions, Some(build_context)))
1450        } else {
1451            // If we're not producing a builtin, we're producing a shared library
1452            // extension module. We currently only support extension modules that
1453            // already have a shared library present. So we simply call into
1454            // the resources collector.
1455            let location = if prefer_in_memory && can_load_dynamic_library_memory {
1456                ConcreteResourceLocation::InMemory
1457            } else {
1458                match relative_path {
1459                    Some(prefix) => ConcreteResourceLocation::RelativePath(prefix),
1460                    None => ConcreteResourceLocation::InMemory,
1461                }
1462            };
1463
1464            let actions = self.add_python_extension_module(extension_module, &location)?;
1465
1466            Ok((actions, None))
1467        }
1468    }
1469
1470    /// Add a built-in extension module.
1471    ///
1472    /// Built-in extension modules are statically linked into the binary and
1473    /// cannot have their location defined.
1474    pub fn add_builtin_python_extension_module(
1475        &mut self,
1476        module: &PythonExtensionModule,
1477    ) -> Result<Vec<AddResourceAction>> {
1478        let entry = self
1479            .resources
1480            .entry(module.name.clone())
1481            .or_insert_with(|| PrePackagedResource {
1482                name: module.name.clone(),
1483                ..PrePackagedResource::default()
1484            });
1485
1486        entry.is_builtin_extension_module = true;
1487        entry.is_package = module.is_package;
1488
1489        Ok(vec![AddResourceAction::AddedBuiltinExtensionModule(
1490            module.name.clone(),
1491        )])
1492    }
1493
1494    /// Add a Python extension module shared library that should be imported from memory.
1495    pub fn add_python_extension_module(
1496        &mut self,
1497        module: &PythonExtensionModule,
1498        location: &ConcreteResourceLocation,
1499    ) -> Result<Vec<AddResourceAction>> {
1500        self.check_policy(location.into())?;
1501
1502        let data = match &module.shared_library {
1503            Some(location) => location.resolve_content()?,
1504            None => return Err(anyhow!("no shared library data present")),
1505        };
1506
1507        match location {
1508            ConcreteResourceLocation::RelativePath(_) => {
1509                if !self
1510                    .allowed_extension_module_locations
1511                    .contains(&AbstractResourceLocation::RelativePath)
1512                {
1513                    return Err(anyhow!("cannot add extension module {} as a file because extension modules as files are not allowed", module.name));
1514                }
1515            }
1516            ConcreteResourceLocation::InMemory => {
1517                if !self
1518                    .allowed_extension_module_locations
1519                    .contains(&AbstractResourceLocation::InMemory)
1520                {
1521                    return Err(anyhow!("cannot add extension module {} for in-memory import because in-memory loading is not supported/allowed", module.name));
1522                }
1523            }
1524        }
1525
1526        let mut depends = vec![];
1527        let mut actions = vec![];
1528
1529        for link in &module.link_libraries {
1530            if link.dynamic_library.is_some() {
1531                let library_location = match location {
1532                    ConcreteResourceLocation::InMemory => ConcreteResourceLocation::InMemory,
1533                    ConcreteResourceLocation::RelativePath(prefix) => {
1534                        // We place the shared library next to the extension module.
1535                        let path = module
1536                            .resolve_path(prefix)
1537                            .parent()
1538                            .ok_or_else(|| anyhow!("unable to resolve parent directory"))?
1539                            .to_path_buf();
1540
1541                        ConcreteResourceLocation::RelativePath(
1542                            path.display().to_string().replace('\\', "/"),
1543                        )
1544                    }
1545                };
1546
1547                let library = SharedLibrary::try_from(link).map_err(|e| anyhow!(e.to_string()))?;
1548
1549                actions.extend(self.add_shared_library(&library, &library_location)?);
1550                depends.push(link.name.to_string());
1551            }
1552        }
1553
1554        let entry = self
1555            .resources
1556            .entry(module.name.to_string())
1557            .or_insert_with(|| PrePackagedResource {
1558                name: module.name.to_string(),
1559                ..PrePackagedResource::default()
1560            });
1561
1562        entry.is_extension_module = true;
1563
1564        if module.is_package {
1565            entry.is_package = true;
1566        }
1567
1568        match location {
1569            ConcreteResourceLocation::InMemory => {
1570                entry.in_memory_extension_module_shared_library = Some(FileData::Memory(data));
1571            }
1572            ConcreteResourceLocation::RelativePath(prefix) => {
1573                entry.relative_path_extension_module_shared_library =
1574                    Some((module.resolve_path(prefix), FileData::Memory(data)));
1575            }
1576        }
1577
1578        entry.shared_library_dependency_names = Some(depends);
1579        actions.push(AddResourceAction::Added(
1580            module.description(),
1581            location.clone(),
1582        ));
1583
1584        Ok(actions)
1585    }
1586
1587    /// Add a shared library to be loaded from a location.
1588    pub fn add_shared_library(
1589        &mut self,
1590        library: &SharedLibrary,
1591        location: &ConcreteResourceLocation,
1592    ) -> Result<Vec<AddResourceAction>> {
1593        self.check_policy(location.into())?;
1594
1595        let entry = self
1596            .resources
1597            .entry(library.name.to_string())
1598            .or_insert_with(|| PrePackagedResource {
1599                name: library.name.to_string(),
1600                ..PrePackagedResource::default()
1601            });
1602
1603        entry.is_shared_library = true;
1604
1605        match location {
1606            ConcreteResourceLocation::InMemory => {
1607                entry.in_memory_shared_library = Some(library.data.clone());
1608            }
1609            ConcreteResourceLocation::RelativePath(prefix) => match &library.filename {
1610                Some(filename) => {
1611                    entry.relative_path_shared_library =
1612                        Some((prefix.to_string(), filename.clone(), library.data.clone()));
1613                }
1614                None => return Err(anyhow!("cannot add shared library without known filename")),
1615            },
1616        }
1617
1618        Ok(vec![AddResourceAction::Added(
1619            library.description(),
1620            location.clone(),
1621        )])
1622    }
1623
1624    pub fn add_file_data(
1625        &mut self,
1626        file: &File,
1627        location: &ConcreteResourceLocation,
1628    ) -> Result<Vec<AddResourceAction>> {
1629        if !self.allow_files {
1630            return Err(anyhow!(
1631                "untyped files are now allowed on this resource collector"
1632            ));
1633        }
1634
1635        self.check_policy(location.into())?;
1636
1637        let entry =
1638            self.resources
1639                .entry(file.path_string())
1640                .or_insert_with(|| PrePackagedResource {
1641                    name: file.path_string(),
1642                    ..PrePackagedResource::default()
1643                });
1644
1645        entry.is_utf8_filename_data = true;
1646        entry.file_executable = file.entry().is_executable();
1647
1648        match location {
1649            ConcreteResourceLocation::InMemory => {
1650                entry.file_data_embedded = Some(file.entry().file_data().clone());
1651            }
1652            ConcreteResourceLocation::RelativePath(prefix) => {
1653                let path = PathBuf::from(prefix).join(file.path());
1654
1655                entry.file_data_utf8_relative_path = Some((
1656                    PathBuf::from(path.display().to_string().replace('\\', "/")),
1657                    file.entry().file_data().clone(),
1658                ));
1659            }
1660        }
1661
1662        Ok(vec![AddResourceAction::Added(
1663            format!("file {}", file.path_string()),
1664            location.clone(),
1665        )])
1666    }
1667
1668    pub fn add_file_data_with_context(
1669        &mut self,
1670        file: &File,
1671        add_context: &PythonResourceAddCollectionContext,
1672    ) -> Result<Vec<AddResourceAction>> {
1673        if !add_context.include {
1674            return Ok(vec![AddResourceAction::NoInclude(format!(
1675                "file {}",
1676                file.path_string()
1677            ))]);
1678        }
1679
1680        self.add_python_resource_with_locations(
1681            &file.into(),
1682            &add_context.location,
1683            &add_context.location_fallback,
1684        )
1685    }
1686
1687    fn add_python_resource_with_locations(
1688        &mut self,
1689        resource: &PythonResource,
1690        location: &ConcreteResourceLocation,
1691        fallback_location: &Option<ConcreteResourceLocation>,
1692    ) -> Result<Vec<AddResourceAction>> {
1693        match resource {
1694            PythonResource::ModuleSource(module) => {
1695                match self
1696                    .add_python_module_source(module, location)
1697                    .with_context(|| format!("adding PythonModuleSource<{}>", module.name))
1698                {
1699                    Ok(actions) => Ok(actions),
1700                    Err(err) => {
1701                        if let Some(location) = fallback_location {
1702                            self.add_python_module_source(module, location)
1703                        } else {
1704                            Err(err)
1705                        }
1706                    }
1707                }
1708            }
1709            PythonResource::ModuleBytecodeRequest(module) => {
1710                match self
1711                    .add_python_module_bytecode_from_source(module, location)
1712                    .with_context(|| {
1713                        format!("adding PythonModuleBytecodeFromSource<{}>", module.name)
1714                    }) {
1715                    Ok(actions) => Ok(actions),
1716                    Err(err) => {
1717                        if let Some(location) = fallback_location {
1718                            self.add_python_module_bytecode_from_source(module, location)
1719                        } else {
1720                            Err(err)
1721                        }
1722                    }
1723                }
1724            }
1725            PythonResource::ModuleBytecode(module) => {
1726                match self
1727                    .add_python_module_bytecode(module, location)
1728                    .with_context(|| format!("adding PythonModuleBytecode<{}>", module.name))
1729                {
1730                    Ok(actions) => Ok(actions),
1731                    Err(err) => {
1732                        if let Some(location) = fallback_location {
1733                            self.add_python_module_bytecode(module, location)
1734                        } else {
1735                            Err(err)
1736                        }
1737                    }
1738                }
1739            }
1740            PythonResource::PackageResource(resource) => {
1741                match self
1742                    .add_python_package_resource(resource, location)
1743                    .with_context(|| {
1744                        format!(
1745                            "adding PythonPackageResource<{}, {}>",
1746                            resource.leaf_package, resource.relative_name
1747                        )
1748                    }) {
1749                    Ok(actions) => Ok(actions),
1750                    Err(err) => {
1751                        if let Some(location) = fallback_location {
1752                            self.add_python_package_resource(resource, location)
1753                        } else {
1754                            Err(err)
1755                        }
1756                    }
1757                }
1758            }
1759            PythonResource::PackageDistributionResource(resource) => {
1760                match self
1761                    .add_python_package_distribution_resource(resource, location)
1762                    .with_context(|| {
1763                        format!(
1764                            "adding PythonPackageDistributionResource<{}, {}>",
1765                            resource.package, resource.name
1766                        )
1767                    }) {
1768                    Ok(actions) => Ok(actions),
1769                    Err(err) => {
1770                        if let Some(location) = fallback_location {
1771                            self.add_python_package_distribution_resource(resource, location)
1772                        } else {
1773                            Err(err)
1774                        }
1775                    }
1776                }
1777            }
1778            PythonResource::File(file) => match self
1779                .add_file_data(file, location)
1780                .with_context(|| format!("adding File<{}>", file.path().display()))
1781            {
1782                Ok(actions) => Ok(actions),
1783                Err(err) => {
1784                    if let Some(location) = fallback_location {
1785                        self.add_file_data(file, location)
1786                    } else {
1787                        Err(err)
1788                    }
1789                }
1790            },
1791            _ => Err(anyhow!("PythonResource variant not yet supported")),
1792        }
1793    }
1794
1795    /// Searches for Python sources for references to __file__.
1796    ///
1797    /// __file__ usage can be problematic for in-memory modules. This method searches
1798    /// for its occurrences and returns module names having it present.
1799    pub fn find_dunder_file(&self) -> Result<BTreeSet<String>> {
1800        let mut res = BTreeSet::new();
1801
1802        for (name, module) in &self.resources {
1803            if let Some(location) = &module.in_memory_source {
1804                if has_dunder_file(&location.resolve_content()?)? {
1805                    res.insert(name.clone());
1806                }
1807            }
1808
1809            if let Some(PythonModuleBytecodeProvider::FromSource(location)) =
1810                &module.in_memory_bytecode
1811            {
1812                if has_dunder_file(&location.resolve_content()?)? {
1813                    res.insert(name.clone());
1814                }
1815            }
1816
1817            if let Some(PythonModuleBytecodeProvider::FromSource(location)) =
1818                &module.in_memory_bytecode_opt1
1819            {
1820                if has_dunder_file(&location.resolve_content()?)? {
1821                    res.insert(name.clone());
1822                }
1823            }
1824
1825            if let Some(PythonModuleBytecodeProvider::FromSource(location)) =
1826                &module.in_memory_bytecode_opt2
1827            {
1828                if has_dunder_file(&location.resolve_content()?)? {
1829                    res.insert(name.clone());
1830                }
1831            }
1832        }
1833
1834        Ok(res)
1835    }
1836
1837    /// Compiles resources into a finalized collection.
1838    ///
1839    /// This will take all resources collected so far and convert them into
1840    /// a collection of `Resource` plus extra file install rules.
1841    ///
1842    /// Missing parent packages will be added automatically.
1843    pub fn compile_resources(
1844        &self,
1845        compiler: &mut dyn PythonBytecodeCompiler,
1846    ) -> Result<CompiledResourcesCollection> {
1847        let mut input_resources = self.resources.clone();
1848        populate_parent_packages(&mut input_resources).context("populating parent packages")?;
1849
1850        let mut resources = BTreeMap::new();
1851        let mut extra_files = Vec::new();
1852
1853        for (name, resource) in &input_resources {
1854            let (entry, installs) = resource
1855                .to_resource(compiler)
1856                .with_context(|| format!("converting {} to resource", name))?;
1857
1858            for install in installs {
1859                extra_files.push(install);
1860            }
1861
1862            resources.insert(name.clone(), entry);
1863        }
1864
1865        Ok(CompiledResourcesCollection {
1866            resources,
1867            extra_files,
1868        })
1869    }
1870}
1871
1872#[cfg(test)]
1873mod tests {
1874    use {
1875        super::*,
1876        crate::{
1877            resource::{LibraryDependency, PythonPackageDistributionResourceFlavor},
1878            testutil::FakeBytecodeCompiler,
1879        },
1880        simple_file_manifest::FileEntry,
1881    };
1882
1883    const DEFAULT_CACHE_TAG: &str = "cpython-39";
1884
1885    #[test]
1886    fn test_resource_conversion_basic() -> Result<()> {
1887        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
1888
1889        let pre = PrePackagedResource {
1890            is_module: true,
1891            name: "module".to_string(),
1892            is_package: true,
1893            is_namespace_package: true,
1894            ..PrePackagedResource::default()
1895        };
1896
1897        let (resource, installs) = pre.to_resource(&mut compiler)?;
1898
1899        assert_eq!(
1900            resource,
1901            Resource {
1902                is_python_module: true,
1903                name: Cow::Owned("module".to_string()),
1904                is_python_package: true,
1905                is_python_namespace_package: true,
1906                ..Resource::default()
1907            }
1908        );
1909
1910        assert!(installs.is_empty());
1911
1912        Ok(())
1913    }
1914
1915    #[test]
1916    fn test_resource_conversion_in_memory_source() -> Result<()> {
1917        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
1918
1919        let pre = PrePackagedResource {
1920            is_module: true,
1921            name: "module".to_string(),
1922            in_memory_source: Some(FileData::Memory(b"source".to_vec())),
1923            ..PrePackagedResource::default()
1924        };
1925
1926        let (resource, installs) = pre.to_resource(&mut compiler)?;
1927
1928        assert_eq!(
1929            resource,
1930            Resource {
1931                is_python_module: true,
1932                name: Cow::Owned("module".to_string()),
1933                in_memory_source: Some(Cow::Owned(b"source".to_vec())),
1934                ..Resource::default()
1935            }
1936        );
1937
1938        assert!(installs.is_empty());
1939
1940        Ok(())
1941    }
1942
1943    #[test]
1944    fn test_resource_conversion_in_memory_bytecode_provided() -> Result<()> {
1945        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
1946
1947        let pre = PrePackagedResource {
1948            is_module: true,
1949            name: "module".to_string(),
1950            in_memory_bytecode: Some(PythonModuleBytecodeProvider::Provided(FileData::Memory(
1951                b"bytecode".to_vec(),
1952            ))),
1953            ..PrePackagedResource::default()
1954        };
1955
1956        let (resource, installs) = pre.to_resource(&mut compiler)?;
1957
1958        assert_eq!(
1959            resource,
1960            Resource {
1961                is_python_module: true,
1962                name: Cow::Owned("module".to_string()),
1963                in_memory_bytecode: Some(Cow::Owned(b"bytecode".to_vec())),
1964                ..Resource::default()
1965            }
1966        );
1967        assert!(installs.is_empty());
1968
1969        Ok(())
1970    }
1971
1972    #[test]
1973    fn test_resource_conversion_in_memory_bytecode_from_source() -> Result<()> {
1974        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
1975
1976        let pre = PrePackagedResource {
1977            is_module: true,
1978            name: "module".to_string(),
1979            in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(FileData::Memory(
1980                b"source".to_vec(),
1981            ))),
1982            ..PrePackagedResource::default()
1983        };
1984
1985        let (resource, installs) = pre.to_resource(&mut compiler)?;
1986
1987        assert_eq!(
1988            resource,
1989            Resource {
1990                is_python_module: true,
1991                name: Cow::Owned("module".to_string()),
1992                in_memory_bytecode: Some(Cow::Owned(b"bc0source".to_vec())),
1993                ..Resource::default()
1994            }
1995        );
1996        assert!(installs.is_empty());
1997
1998        Ok(())
1999    }
2000
2001    #[test]
2002    fn test_resource_conversion_in_memory_bytecode_opt1_provided() -> Result<()> {
2003        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2004
2005        let pre = PrePackagedResource {
2006            is_module: true,
2007            name: "module".to_string(),
2008            in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::Provided(
2009                FileData::Memory(b"bytecode".to_vec()),
2010            )),
2011            ..PrePackagedResource::default()
2012        };
2013
2014        let (resource, installs) = pre.to_resource(&mut compiler)?;
2015
2016        assert_eq!(
2017            resource,
2018            Resource {
2019                is_python_module: true,
2020                name: Cow::Owned("module".to_string()),
2021                in_memory_bytecode_opt1: Some(Cow::Owned(b"bytecode".to_vec())),
2022                ..Resource::default()
2023            }
2024        );
2025        assert!(installs.is_empty());
2026
2027        Ok(())
2028    }
2029
2030    #[test]
2031    fn test_resource_conversion_in_memory_bytecode_opt1_from_source() -> Result<()> {
2032        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2033
2034        let pre = PrePackagedResource {
2035            is_module: true,
2036            name: "module".to_string(),
2037            in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::FromSource(
2038                FileData::Memory(b"source".to_vec()),
2039            )),
2040            ..PrePackagedResource::default()
2041        };
2042
2043        let (resource, installs) = pre.to_resource(&mut compiler)?;
2044
2045        assert_eq!(
2046            resource,
2047            Resource {
2048                is_python_module: true,
2049                name: Cow::Owned("module".to_string()),
2050                in_memory_bytecode_opt1: Some(Cow::Owned(b"bc1source".to_vec())),
2051                ..Resource::default()
2052            }
2053        );
2054        assert!(installs.is_empty());
2055
2056        Ok(())
2057    }
2058
2059    #[test]
2060    fn test_resource_conversion_in_memory_bytecode_opt2_provided() -> Result<()> {
2061        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2062
2063        let pre = PrePackagedResource {
2064            is_module: true,
2065            name: "module".to_string(),
2066            in_memory_bytecode_opt2: Some(PythonModuleBytecodeProvider::Provided(
2067                FileData::Memory(b"bytecode".to_vec()),
2068            )),
2069            ..PrePackagedResource::default()
2070        };
2071
2072        let (resource, installs) = pre.to_resource(&mut compiler)?;
2073
2074        assert_eq!(
2075            resource,
2076            Resource {
2077                is_python_module: true,
2078                name: Cow::Owned("module".to_string()),
2079                in_memory_bytecode_opt2: Some(Cow::Owned(b"bytecode".to_vec())),
2080                ..Resource::default()
2081            }
2082        );
2083        assert!(installs.is_empty());
2084
2085        Ok(())
2086    }
2087
2088    #[test]
2089    fn test_resource_conversion_in_memory_bytecode_opt2_from_source() -> Result<()> {
2090        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2091
2092        let pre = PrePackagedResource {
2093            is_module: true,
2094            name: "module".to_string(),
2095            in_memory_bytecode_opt2: Some(PythonModuleBytecodeProvider::FromSource(
2096                FileData::Memory(b"source".to_vec()),
2097            )),
2098            ..PrePackagedResource::default()
2099        };
2100
2101        let (resource, installs) = pre.to_resource(&mut compiler)?;
2102
2103        assert_eq!(
2104            resource,
2105            Resource {
2106                is_python_module: true,
2107                name: Cow::Owned("module".to_string()),
2108                in_memory_bytecode_opt2: Some(Cow::Owned(b"bc2source".to_vec())),
2109                ..Resource::default()
2110            }
2111        );
2112        assert!(installs.is_empty());
2113
2114        Ok(())
2115    }
2116
2117    #[test]
2118    fn test_resource_conversion_in_memory_extension_module_shared_library() -> Result<()> {
2119        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2120
2121        let pre = PrePackagedResource {
2122            is_module: true,
2123            name: "module".to_string(),
2124            in_memory_extension_module_shared_library: Some(FileData::Memory(b"library".to_vec())),
2125            ..PrePackagedResource::default()
2126        };
2127
2128        let (resource, installs) = pre.to_resource(&mut compiler)?;
2129
2130        assert_eq!(
2131            resource,
2132            Resource {
2133                is_python_module: true,
2134                name: Cow::Owned("module".to_string()),
2135                in_memory_extension_module_shared_library: Some(Cow::Owned(b"library".to_vec())),
2136                ..Resource::default()
2137            }
2138        );
2139
2140        assert!(installs.is_empty());
2141
2142        Ok(())
2143    }
2144
2145    #[test]
2146    fn test_resource_conversion_in_memory_package_resources() -> Result<()> {
2147        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2148
2149        let mut resources = BTreeMap::new();
2150        resources.insert("foo".to_string(), FileData::Memory(b"value".to_vec()));
2151
2152        let pre = PrePackagedResource {
2153            is_module: true,
2154            name: "module".to_string(),
2155            in_memory_resources: Some(resources),
2156            ..PrePackagedResource::default()
2157        };
2158
2159        let (resource, installs) = pre.to_resource(&mut compiler)?;
2160
2161        let mut resources = HashMap::new();
2162        resources.insert(Cow::Owned("foo".to_string()), Cow::Owned(b"value".to_vec()));
2163
2164        assert_eq!(
2165            resource,
2166            Resource {
2167                is_python_module: true,
2168                name: Cow::Owned("module".to_string()),
2169                in_memory_package_resources: Some(resources),
2170                ..Resource::default()
2171            }
2172        );
2173
2174        assert!(installs.is_empty());
2175
2176        Ok(())
2177    }
2178
2179    #[test]
2180    fn test_resource_conversion_in_memory_distribution_resources() -> Result<()> {
2181        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2182
2183        let mut resources = BTreeMap::new();
2184        resources.insert("foo".to_string(), FileData::Memory(b"value".to_vec()));
2185
2186        let pre = PrePackagedResource {
2187            is_module: true,
2188            name: "module".to_string(),
2189            in_memory_distribution_resources: Some(resources),
2190            ..PrePackagedResource::default()
2191        };
2192
2193        let (resource, installs) = pre.to_resource(&mut compiler)?;
2194
2195        let mut resources = HashMap::new();
2196        resources.insert(Cow::Owned("foo".to_string()), Cow::Owned(b"value".to_vec()));
2197
2198        assert_eq!(
2199            resource,
2200            Resource {
2201                is_python_module: true,
2202                name: Cow::Owned("module".to_string()),
2203                in_memory_distribution_resources: Some(resources),
2204                ..Resource::default()
2205            }
2206        );
2207
2208        assert!(installs.is_empty());
2209
2210        Ok(())
2211    }
2212
2213    #[test]
2214    fn test_resource_conversion_in_memory_shared_library() -> Result<()> {
2215        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2216
2217        let pre = PrePackagedResource {
2218            is_shared_library: true,
2219            name: "module".to_string(),
2220            in_memory_shared_library: Some(FileData::Memory(b"library".to_vec())),
2221            shared_library_dependency_names: Some(vec!["foo".to_string(), "bar".to_string()]),
2222            ..PrePackagedResource::default()
2223        };
2224
2225        let (resource, installs) = pre.to_resource(&mut compiler)?;
2226
2227        assert_eq!(
2228            resource,
2229            Resource {
2230                is_shared_library: true,
2231                name: Cow::Owned("module".to_string()),
2232                in_memory_shared_library: Some(Cow::Owned(b"library".to_vec())),
2233                shared_library_dependency_names: Some(vec![
2234                    Cow::Owned("foo".to_string()),
2235                    Cow::Owned("bar".to_string())
2236                ]),
2237                ..Resource::default()
2238            }
2239        );
2240
2241        assert!(installs.is_empty());
2242
2243        Ok(())
2244    }
2245
2246    #[test]
2247    fn test_resource_conversion_relative_path_module_source() -> Result<()> {
2248        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2249
2250        let pre = PrePackagedResource {
2251            is_module: true,
2252            name: "module".to_string(),
2253            relative_path_module_source: Some((
2254                "prefix".to_string(),
2255                FileData::Memory(b"source".to_vec()),
2256            )),
2257            ..PrePackagedResource::default()
2258        };
2259
2260        let (resource, installs) = pre.to_resource(&mut compiler)?;
2261
2262        assert_eq!(
2263            resource,
2264            Resource {
2265                is_python_module: true,
2266                name: Cow::Owned("module".to_string()),
2267                relative_path_module_source: Some(Cow::Owned(PathBuf::from("prefix/module.py"))),
2268                ..Resource::default()
2269            }
2270        );
2271
2272        assert_eq!(
2273            installs,
2274            vec![(
2275                PathBuf::from("prefix/module.py"),
2276                FileData::Memory(b"source".to_vec()),
2277                false
2278            )]
2279        );
2280
2281        Ok(())
2282    }
2283
2284    #[test]
2285    fn test_resource_conversion_relative_path_module_bytecode_provided() -> Result<()> {
2286        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2287
2288        let pre = PrePackagedResource {
2289            is_module: true,
2290            name: "foo.bar".to_string(),
2291            relative_path_bytecode: Some((
2292                "prefix".to_string(),
2293                "tag".to_string(),
2294                PythonModuleBytecodeProvider::Provided(FileData::Memory(b"bytecode".to_vec())),
2295            )),
2296            ..PrePackagedResource::default()
2297        };
2298
2299        let (resource, installs) = pre.to_resource(&mut compiler)?;
2300
2301        assert_eq!(
2302            resource,
2303            Resource {
2304                is_python_module: true,
2305                name: Cow::Owned("foo.bar".to_string()),
2306                relative_path_module_bytecode: Some(Cow::Owned(PathBuf::from(
2307                    "prefix/foo/__pycache__/bar.tag.pyc"
2308                ))),
2309                ..Resource::default()
2310            }
2311        );
2312
2313        assert_eq!(
2314            installs,
2315            vec![(
2316                PathBuf::from("prefix/foo/__pycache__/bar.tag.pyc"),
2317                FileData::Memory(
2318                    b"\x2a\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00bytecode"
2319                        .to_vec()
2320                ),
2321                false
2322            )]
2323        );
2324
2325        Ok(())
2326    }
2327
2328    #[test]
2329    fn test_resource_conversion_relative_path_module_bytecode_from_source() -> Result<()> {
2330        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2331
2332        let pre = PrePackagedResource {
2333            is_module: true,
2334            name: "foo.bar".to_string(),
2335            relative_path_bytecode: Some((
2336                "prefix".to_string(),
2337                "tag".to_string(),
2338                PythonModuleBytecodeProvider::FromSource(FileData::Memory(b"source".to_vec())),
2339            )),
2340            ..PrePackagedResource::default()
2341        };
2342
2343        let (resource, installs) = pre.to_resource(&mut compiler)?;
2344
2345        assert_eq!(
2346            resource,
2347            Resource {
2348                is_python_module: true,
2349                name: Cow::Owned("foo.bar".to_string()),
2350                relative_path_module_bytecode: Some(Cow::Owned(PathBuf::from(
2351                    "prefix/foo/__pycache__/bar.tag.pyc"
2352                ))),
2353                ..Resource::default()
2354            }
2355        );
2356
2357        assert_eq!(
2358            installs,
2359            vec![(
2360                PathBuf::from("prefix/foo/__pycache__/bar.tag.pyc"),
2361                FileData::Memory(b"bc0source".to_vec()),
2362                false
2363            )]
2364        );
2365
2366        Ok(())
2367    }
2368
2369    #[test]
2370    fn test_resource_conversion_relative_path_module_bytecode_opt1_provided() -> Result<()> {
2371        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2372
2373        let pre = PrePackagedResource {
2374            is_module: true,
2375            name: "foo.bar".to_string(),
2376            relative_path_bytecode_opt1: Some((
2377                "prefix".to_string(),
2378                "tag".to_string(),
2379                PythonModuleBytecodeProvider::Provided(FileData::Memory(b"bytecode".to_vec())),
2380            )),
2381            ..PrePackagedResource::default()
2382        };
2383
2384        let (resource, installs) = pre.to_resource(&mut compiler)?;
2385
2386        assert_eq!(
2387            resource,
2388            Resource {
2389                is_python_module: true,
2390                name: Cow::Owned("foo.bar".to_string()),
2391                relative_path_module_bytecode_opt1: Some(Cow::Owned(PathBuf::from(
2392                    "prefix/foo/__pycache__/bar.tag.opt-1.pyc"
2393                ))),
2394                ..Resource::default()
2395            }
2396        );
2397
2398        assert_eq!(
2399            installs,
2400            vec![(
2401                PathBuf::from("prefix/foo/__pycache__/bar.tag.opt-1.pyc"),
2402                FileData::Memory(
2403                    b"\x2a\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00bytecode"
2404                        .to_vec()
2405                ),
2406                false
2407            )]
2408        );
2409
2410        Ok(())
2411    }
2412
2413    #[test]
2414    fn test_resource_conversion_relative_path_module_bytecode_opt1_from_source() -> Result<()> {
2415        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2416
2417        let pre = PrePackagedResource {
2418            is_module: true,
2419            name: "foo.bar".to_string(),
2420            relative_path_bytecode_opt1: Some((
2421                "prefix".to_string(),
2422                "tag".to_string(),
2423                PythonModuleBytecodeProvider::FromSource(FileData::Memory(b"source".to_vec())),
2424            )),
2425            ..PrePackagedResource::default()
2426        };
2427
2428        let (resource, installs) = pre.to_resource(&mut compiler)?;
2429
2430        assert_eq!(
2431            resource,
2432            Resource {
2433                is_python_module: true,
2434                name: Cow::Owned("foo.bar".to_string()),
2435                relative_path_module_bytecode_opt1: Some(Cow::Owned(PathBuf::from(
2436                    "prefix/foo/__pycache__/bar.tag.opt-1.pyc"
2437                ))),
2438                ..Resource::default()
2439            }
2440        );
2441
2442        assert_eq!(
2443            installs,
2444            vec![(
2445                PathBuf::from("prefix/foo/__pycache__/bar.tag.opt-1.pyc"),
2446                FileData::Memory(b"bc1source".to_vec()),
2447                false
2448            )]
2449        );
2450
2451        Ok(())
2452    }
2453
2454    #[test]
2455    fn test_resource_conversion_relative_path_module_bytecode_opt2_provided() -> Result<()> {
2456        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2457
2458        let pre = PrePackagedResource {
2459            is_module: true,
2460            name: "foo.bar".to_string(),
2461            relative_path_bytecode_opt2: Some((
2462                "prefix".to_string(),
2463                "tag".to_string(),
2464                PythonModuleBytecodeProvider::Provided(FileData::Memory(b"bytecode".to_vec())),
2465            )),
2466            ..PrePackagedResource::default()
2467        };
2468
2469        let (resource, installs) = pre.to_resource(&mut compiler)?;
2470
2471        assert_eq!(
2472            resource,
2473            Resource {
2474                is_python_module: true,
2475                name: Cow::Owned("foo.bar".to_string()),
2476                relative_path_module_bytecode_opt2: Some(Cow::Owned(PathBuf::from(
2477                    "prefix/foo/__pycache__/bar.tag.opt-2.pyc"
2478                ))),
2479                ..Resource::default()
2480            }
2481        );
2482
2483        assert_eq!(
2484            installs,
2485            vec![(
2486                PathBuf::from("prefix/foo/__pycache__/bar.tag.opt-2.pyc"),
2487                FileData::Memory(
2488                    b"\x2a\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00bytecode"
2489                        .to_vec()
2490                ),
2491                false
2492            )]
2493        );
2494
2495        Ok(())
2496    }
2497
2498    #[test]
2499    fn test_resource_conversion_relative_path_module_bytecode_opt2_from_source() -> Result<()> {
2500        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2501
2502        let pre = PrePackagedResource {
2503            is_module: true,
2504            name: "foo.bar".to_string(),
2505            relative_path_bytecode_opt2: Some((
2506                "prefix".to_string(),
2507                "tag".to_string(),
2508                PythonModuleBytecodeProvider::FromSource(FileData::Memory(b"source".to_vec())),
2509            )),
2510            ..PrePackagedResource::default()
2511        };
2512
2513        let (resource, installs) = pre.to_resource(&mut compiler)?;
2514
2515        assert_eq!(
2516            resource,
2517            Resource {
2518                is_python_module: true,
2519                name: Cow::Owned("foo.bar".to_string()),
2520                relative_path_module_bytecode_opt2: Some(Cow::Owned(PathBuf::from(
2521                    "prefix/foo/__pycache__/bar.tag.opt-2.pyc"
2522                ))),
2523                ..Resource::default()
2524            }
2525        );
2526
2527        assert_eq!(
2528            installs,
2529            vec![(
2530                PathBuf::from("prefix/foo/__pycache__/bar.tag.opt-2.pyc"),
2531                FileData::Memory(b"bc2source".to_vec()),
2532                false
2533            )]
2534        );
2535
2536        Ok(())
2537    }
2538
2539    #[test]
2540    fn test_resource_conversion_relative_path_extension_module_shared_library() -> Result<()> {
2541        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2542
2543        let pre = PrePackagedResource {
2544            is_module: true,
2545            name: "module".to_string(),
2546            relative_path_extension_module_shared_library: Some((
2547                PathBuf::from("prefix/ext.so"),
2548                FileData::Memory(b"data".to_vec()),
2549            )),
2550            ..PrePackagedResource::default()
2551        };
2552
2553        let (resource, installs) = pre.to_resource(&mut compiler)?;
2554
2555        assert_eq!(
2556            resource,
2557            Resource {
2558                is_python_module: true,
2559                name: Cow::Owned("module".to_string()),
2560                relative_path_extension_module_shared_library: Some(Cow::Owned(PathBuf::from(
2561                    "prefix/ext.so"
2562                ))),
2563                ..Resource::default()
2564            }
2565        );
2566
2567        assert_eq!(
2568            installs,
2569            vec![(
2570                PathBuf::from("prefix/ext.so"),
2571                FileData::Memory(b"data".to_vec()),
2572                true
2573            )]
2574        );
2575
2576        Ok(())
2577    }
2578
2579    #[test]
2580    fn test_resource_conversion_relative_path_package_resources() -> Result<()> {
2581        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2582
2583        let mut resources = BTreeMap::new();
2584        resources.insert(
2585            "foo.txt".to_string(),
2586            (
2587                PathBuf::from("module/foo.txt"),
2588                FileData::Memory(b"data".to_vec()),
2589            ),
2590        );
2591        resources.insert(
2592            "bar.txt".to_string(),
2593            (
2594                PathBuf::from("module/bar.txt"),
2595                FileData::Memory(b"bar".to_vec()),
2596            ),
2597        );
2598
2599        let pre = PrePackagedResource {
2600            is_module: true,
2601            name: "module".to_string(),
2602            is_package: true,
2603            relative_path_package_resources: Some(resources),
2604            ..PrePackagedResource::default()
2605        };
2606
2607        let (resource, installs) = pre.to_resource(&mut compiler)?;
2608
2609        let mut resources = HashMap::new();
2610        resources.insert(
2611            Cow::Owned("foo.txt".to_string()),
2612            Cow::Owned(PathBuf::from("module/foo.txt")),
2613        );
2614        resources.insert(
2615            Cow::Owned("bar.txt".to_string()),
2616            Cow::Owned(PathBuf::from("module/bar.txt")),
2617        );
2618
2619        assert_eq!(
2620            resource,
2621            Resource {
2622                is_python_module: true,
2623                name: Cow::Owned("module".to_string()),
2624                is_python_package: true,
2625                relative_path_package_resources: Some(resources),
2626                ..Resource::default()
2627            }
2628        );
2629
2630        assert_eq!(
2631            installs,
2632            vec![
2633                (
2634                    PathBuf::from("module/bar.txt"),
2635                    FileData::Memory(b"bar".to_vec()),
2636                    false
2637                ),
2638                (
2639                    PathBuf::from("module/foo.txt"),
2640                    FileData::Memory(b"data".to_vec()),
2641                    false
2642                ),
2643            ]
2644        );
2645
2646        Ok(())
2647    }
2648
2649    #[test]
2650    fn test_resource_conversion_relative_path_distribution_resources() -> Result<()> {
2651        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2652
2653        let mut resources = BTreeMap::new();
2654        resources.insert(
2655            "foo.txt".to_string(),
2656            (PathBuf::from("foo.txt"), FileData::Memory(b"data".to_vec())),
2657        );
2658
2659        let pre = PrePackagedResource {
2660            is_module: true,
2661            name: "module".to_string(),
2662            relative_path_distribution_resources: Some(resources),
2663            ..PrePackagedResource::default()
2664        };
2665
2666        let (resource, installs) = pre.to_resource(&mut compiler)?;
2667
2668        let mut resources = HashMap::new();
2669        resources.insert(
2670            Cow::Owned("foo.txt".to_string()),
2671            Cow::Owned(PathBuf::from("foo.txt")),
2672        );
2673
2674        assert_eq!(
2675            resource,
2676            Resource {
2677                is_python_module: true,
2678                name: Cow::Owned("module".to_string()),
2679                relative_path_distribution_resources: Some(resources),
2680                ..Resource::default()
2681            }
2682        );
2683
2684        assert_eq!(
2685            installs,
2686            vec![(
2687                PathBuf::from("foo.txt"),
2688                FileData::Memory(b"data".to_vec()),
2689                false
2690            )]
2691        );
2692
2693        Ok(())
2694    }
2695
2696    #[test]
2697    fn test_resource_conversion_relative_path_shared_library() -> Result<()> {
2698        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2699
2700        let pre = PrePackagedResource {
2701            is_shared_library: true,
2702            name: "libfoo".to_string(),
2703            relative_path_shared_library: Some((
2704                "prefix".to_string(),
2705                PathBuf::from("libfoo.so"),
2706                FileData::Memory(b"data".to_vec()),
2707            )),
2708            ..PrePackagedResource::default()
2709        };
2710
2711        let (resource, installs) = pre.to_resource(&mut compiler)?;
2712
2713        assert_eq!(
2714            resource,
2715            Resource {
2716                is_shared_library: true,
2717                name: Cow::Owned("libfoo".to_string()),
2718                ..Resource::default()
2719            }
2720        );
2721
2722        assert_eq!(
2723            installs,
2724            vec![(
2725                PathBuf::from("prefix/libfoo.so"),
2726                FileData::Memory(b"data".to_vec()),
2727                true
2728            )]
2729        );
2730
2731        Ok(())
2732    }
2733
2734    #[test]
2735    fn test_populate_parent_packages_in_memory_source() -> Result<()> {
2736        let mut h = BTreeMap::new();
2737        h.insert(
2738            "root.parent.child".to_string(),
2739            PrePackagedResource {
2740                is_module: true,
2741                name: "root.parent.child".to_string(),
2742                in_memory_source: Some(FileData::Memory(vec![42])),
2743                is_package: true,
2744                ..PrePackagedResource::default()
2745            },
2746        );
2747
2748        populate_parent_packages(&mut h)?;
2749
2750        assert_eq!(h.len(), 3);
2751        assert_eq!(
2752            h.get("root.parent"),
2753            Some(&PrePackagedResource {
2754                is_module: true,
2755                name: "root.parent".to_string(),
2756                is_package: true,
2757                in_memory_source: Some(FileData::Memory(vec![])),
2758                ..PrePackagedResource::default()
2759            })
2760        );
2761        assert_eq!(
2762            h.get("root"),
2763            Some(&PrePackagedResource {
2764                is_module: true,
2765                name: "root".to_string(),
2766                is_package: true,
2767                in_memory_source: Some(FileData::Memory(vec![])),
2768                ..PrePackagedResource::default()
2769            })
2770        );
2771
2772        Ok(())
2773    }
2774
2775    #[test]
2776    fn test_populate_parent_packages_relative_path_source() -> Result<()> {
2777        let mut h = BTreeMap::new();
2778        h.insert(
2779            "root.parent.child".to_string(),
2780            PrePackagedResource {
2781                is_module: true,
2782                name: "root.parent.child".to_string(),
2783                relative_path_module_source: Some((
2784                    "prefix".to_string(),
2785                    FileData::Memory(vec![42]),
2786                )),
2787                is_package: true,
2788                ..PrePackagedResource::default()
2789            },
2790        );
2791
2792        populate_parent_packages(&mut h)?;
2793
2794        assert_eq!(h.len(), 3);
2795        assert_eq!(
2796            h.get("root.parent"),
2797            Some(&PrePackagedResource {
2798                is_module: true,
2799                name: "root.parent".to_string(),
2800                is_package: true,
2801                relative_path_module_source: Some(("prefix".to_string(), FileData::Memory(vec![]))),
2802                ..PrePackagedResource::default()
2803            })
2804        );
2805        assert_eq!(
2806            h.get("root"),
2807            Some(&PrePackagedResource {
2808                is_module: true,
2809                name: "root".to_string(),
2810                is_package: true,
2811                relative_path_module_source: Some(("prefix".to_string(), FileData::Memory(vec![]))),
2812                ..PrePackagedResource::default()
2813            })
2814        );
2815
2816        Ok(())
2817    }
2818
2819    #[test]
2820    fn test_populate_parent_packages_in_memory_bytecode() -> Result<()> {
2821        let mut h = BTreeMap::new();
2822        h.insert(
2823            "root.parent.child".to_string(),
2824            PrePackagedResource {
2825                is_module: true,
2826                name: "root.parent.child".to_string(),
2827                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
2828                    FileData::Memory(vec![42]),
2829                )),
2830                is_package: true,
2831                ..PrePackagedResource::default()
2832            },
2833        );
2834
2835        populate_parent_packages(&mut h)?;
2836
2837        assert_eq!(h.len(), 3);
2838        assert_eq!(
2839            h.get("root.parent"),
2840            Some(&PrePackagedResource {
2841                is_module: true,
2842                name: "root.parent".to_string(),
2843                is_package: true,
2844                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
2845                    FileData::Memory(vec![])
2846                )),
2847                ..PrePackagedResource::default()
2848            })
2849        );
2850        assert_eq!(
2851            h.get("root"),
2852            Some(&PrePackagedResource {
2853                is_module: true,
2854                name: "root".to_string(),
2855                is_package: true,
2856                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
2857                    FileData::Memory(vec![])
2858                )),
2859                ..PrePackagedResource::default()
2860            })
2861        );
2862
2863        Ok(())
2864    }
2865
2866    #[test]
2867    fn test_populate_parent_packages_distribution_extension_module() -> Result<()> {
2868        let mut h = BTreeMap::new();
2869        h.insert(
2870            "foo.bar".to_string(),
2871            PrePackagedResource {
2872                is_builtin_extension_module: true,
2873                name: "foo.bar".to_string(),
2874                relative_path_extension_module_shared_library: Some((
2875                    PathBuf::from("prefix/foo/bar.so"),
2876                    FileData::Memory(vec![42]),
2877                )),
2878                ..PrePackagedResource::default()
2879            },
2880        );
2881
2882        populate_parent_packages(&mut h)?;
2883
2884        assert_eq!(
2885            h.get("foo"),
2886            Some(&PrePackagedResource {
2887                is_module: true,
2888                name: "foo".to_string(),
2889                is_package: true,
2890                ..PrePackagedResource::default()
2891            })
2892        );
2893
2894        Ok(())
2895    }
2896
2897    #[test]
2898    fn test_populate_parent_packages_relative_extension_module() -> Result<()> {
2899        let mut h = BTreeMap::new();
2900        h.insert(
2901            "foo.bar".to_string(),
2902            PrePackagedResource {
2903                is_extension_module: true,
2904                name: "foo.bar".to_string(),
2905                relative_path_extension_module_shared_library: Some((
2906                    PathBuf::from("prefix/foo/bar.so"),
2907                    FileData::Memory(vec![42]),
2908                )),
2909                ..PrePackagedResource::default()
2910            },
2911        );
2912
2913        populate_parent_packages(&mut h)?;
2914
2915        assert_eq!(h.len(), 2);
2916
2917        assert_eq!(
2918            h.get("foo"),
2919            Some(&PrePackagedResource {
2920                is_module: true,
2921                name: "foo".to_string(),
2922                is_package: true,
2923                ..PrePackagedResource::default()
2924            })
2925        );
2926
2927        Ok(())
2928    }
2929
2930    #[test]
2931    fn test_add_in_memory_source_module() -> Result<()> {
2932        let mut r = PythonResourceCollector::new(
2933            vec![AbstractResourceLocation::InMemory],
2934            vec![],
2935            false,
2936            false,
2937        );
2938        r.add_python_module_source(
2939            &PythonModuleSource {
2940                name: "foo".to_string(),
2941                source: FileData::Memory(vec![42]),
2942                is_package: false,
2943                cache_tag: DEFAULT_CACHE_TAG.to_string(),
2944                is_stdlib: false,
2945                is_test: false,
2946            },
2947            &ConcreteResourceLocation::InMemory,
2948        )?;
2949
2950        assert!(r.resources.contains_key("foo"));
2951        assert_eq!(
2952            r.resources.get("foo"),
2953            Some(&PrePackagedResource {
2954                is_module: true,
2955                name: "foo".to_string(),
2956                is_package: false,
2957                in_memory_source: Some(FileData::Memory(vec![42])),
2958                ..PrePackagedResource::default()
2959            })
2960        );
2961
2962        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
2963
2964        let resources = r.compile_resources(&mut compiler)?;
2965
2966        assert_eq!(resources.resources.len(), 1);
2967        assert_eq!(
2968            resources.resources.get("foo"),
2969            Some(&Resource {
2970                is_python_module: true,
2971                name: Cow::Owned("foo".to_string()),
2972                in_memory_source: Some(Cow::Owned(vec![42])),
2973                ..Resource::default()
2974            })
2975        );
2976        assert!(resources.extra_files.is_empty());
2977
2978        Ok(())
2979    }
2980
2981    #[test]
2982    fn test_add_in_memory_source_module_parents() -> Result<()> {
2983        let mut r = PythonResourceCollector::new(
2984            vec![AbstractResourceLocation::InMemory],
2985            vec![],
2986            false,
2987            false,
2988        );
2989        r.add_python_module_source(
2990            &PythonModuleSource {
2991                name: "root.parent.child".to_string(),
2992                source: FileData::Memory(vec![42]),
2993                is_package: true,
2994                cache_tag: DEFAULT_CACHE_TAG.to_string(),
2995                is_stdlib: false,
2996                is_test: false,
2997            },
2998            &ConcreteResourceLocation::InMemory,
2999        )?;
3000
3001        assert_eq!(r.resources.len(), 1);
3002        assert_eq!(
3003            r.resources.get("root.parent.child"),
3004            Some(&PrePackagedResource {
3005                is_module: true,
3006                name: "root.parent.child".to_string(),
3007                is_package: true,
3008                in_memory_source: Some(FileData::Memory(vec![42])),
3009                ..PrePackagedResource::default()
3010            })
3011        );
3012
3013        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3014
3015        let resources = r.compile_resources(&mut compiler)?;
3016
3017        assert_eq!(resources.resources.len(), 3);
3018        assert_eq!(
3019            resources.resources.get("root"),
3020            Some(&Resource {
3021                is_python_module: true,
3022                name: Cow::Owned("root".to_string()),
3023                is_python_package: true,
3024                in_memory_source: Some(Cow::Owned(vec![])),
3025                ..Resource::default()
3026            })
3027        );
3028        assert_eq!(
3029            resources.resources.get("root.parent"),
3030            Some(&Resource {
3031                is_python_module: true,
3032                name: Cow::Owned("root.parent".to_string()),
3033                is_python_package: true,
3034                in_memory_source: Some(Cow::Owned(vec![])),
3035                ..Resource::default()
3036            })
3037        );
3038        assert_eq!(
3039            resources.resources.get("root.parent.child"),
3040            Some(&Resource {
3041                is_python_module: true,
3042                name: Cow::Owned("root.parent.child".to_string()),
3043                is_python_package: true,
3044                in_memory_source: Some(Cow::Owned(vec![42])),
3045                ..Resource::default()
3046            })
3047        );
3048
3049        assert!(resources.extra_files.is_empty());
3050
3051        Ok(())
3052    }
3053
3054    #[test]
3055    fn test_add_relative_path_source_module() -> Result<()> {
3056        let mut r = PythonResourceCollector::new(
3057            vec![AbstractResourceLocation::RelativePath],
3058            vec![],
3059            false,
3060            false,
3061        );
3062        r.add_python_module_source(
3063            &PythonModuleSource {
3064                name: "foo.bar".to_string(),
3065                source: FileData::Memory(vec![42]),
3066                is_package: false,
3067                cache_tag: DEFAULT_CACHE_TAG.to_string(),
3068                is_stdlib: false,
3069                is_test: false,
3070            },
3071            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
3072        )?;
3073
3074        assert_eq!(r.resources.len(), 1);
3075        assert_eq!(
3076            r.resources.get("foo.bar"),
3077            Some(&PrePackagedResource {
3078                is_module: true,
3079                name: "foo.bar".to_string(),
3080                is_package: false,
3081                relative_path_module_source: Some((
3082                    "prefix".to_string(),
3083                    FileData::Memory(vec![42])
3084                )),
3085                ..PrePackagedResource::default()
3086            })
3087        );
3088
3089        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3090
3091        let resources = r.compile_resources(&mut compiler)?;
3092
3093        assert_eq!(resources.resources.len(), 2);
3094        assert_eq!(
3095            resources.resources.get("foo"),
3096            Some(&Resource {
3097                is_python_module: true,
3098                name: Cow::Owned("foo".to_string()),
3099                is_python_package: true,
3100                relative_path_module_source: Some(Cow::Owned(PathBuf::from(
3101                    "prefix/foo/__init__.py"
3102                ))),
3103                ..Resource::default()
3104            })
3105        );
3106        assert_eq!(
3107            resources.resources.get("foo.bar"),
3108            Some(&Resource {
3109                is_python_module: true,
3110                name: Cow::Owned("foo.bar".to_string()),
3111                relative_path_module_source: Some(Cow::Owned(PathBuf::from("prefix/foo/bar.py"))),
3112                ..Resource::default()
3113            })
3114        );
3115        assert_eq!(
3116            resources.extra_files,
3117            vec![
3118                (
3119                    PathBuf::from("prefix/foo/__init__.py"),
3120                    FileData::Memory(vec![]),
3121                    false
3122                ),
3123                (
3124                    PathBuf::from("prefix/foo/bar.py"),
3125                    FileData::Memory(vec![42]),
3126                    false
3127                )
3128            ]
3129        );
3130
3131        Ok(())
3132    }
3133
3134    #[test]
3135    fn test_add_module_source_with_context() -> Result<()> {
3136        let mut r = PythonResourceCollector::new(
3137            vec![AbstractResourceLocation::InMemory],
3138            vec![],
3139            false,
3140            false,
3141        );
3142
3143        let module = PythonModuleSource {
3144            name: "foo".to_string(),
3145            source: FileData::Memory(vec![42]),
3146            is_package: false,
3147            cache_tag: DEFAULT_CACHE_TAG.to_string(),
3148            is_stdlib: false,
3149            is_test: false,
3150        };
3151
3152        let mut add_context = PythonResourceAddCollectionContext {
3153            include: false,
3154            location: ConcreteResourceLocation::InMemory,
3155            location_fallback: None,
3156            store_source: false,
3157            optimize_level_zero: false,
3158            optimize_level_one: false,
3159            optimize_level_two: false,
3160        };
3161
3162        // include=false is a noop.
3163        assert!(r.resources.is_empty());
3164        r.add_python_module_source_with_context(&module, &add_context)?;
3165        assert!(r.resources.is_empty());
3166
3167        add_context.include = true;
3168
3169        // store_source=false is a noop.
3170        r.add_python_module_source_with_context(&module, &add_context)?;
3171        assert!(r.resources.is_empty());
3172
3173        add_context.store_source = true;
3174
3175        // store_source=true adds just the source.
3176        r.add_python_module_source_with_context(&module, &add_context)?;
3177        assert_eq!(
3178            r.resources.get(&module.name),
3179            Some(&PrePackagedResource {
3180                is_module: true,
3181                name: module.name.clone(),
3182                is_package: module.is_package,
3183                in_memory_source: Some(module.source.clone()),
3184                ..PrePackagedResource::default()
3185            })
3186        );
3187
3188        r.resources.clear();
3189        add_context.store_source = false;
3190
3191        // optimize_level_zero stores the bytecode.
3192
3193        add_context.optimize_level_zero = true;
3194        r.add_python_module_source_with_context(&module, &add_context)?;
3195        assert_eq!(
3196            r.resources.get(&module.name),
3197            Some(&PrePackagedResource {
3198                is_module: true,
3199                name: module.name.clone(),
3200                is_package: module.is_package,
3201                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
3202                    module.source.clone()
3203                )),
3204                ..PrePackagedResource::default()
3205            })
3206        );
3207
3208        r.resources.clear();
3209        add_context.optimize_level_zero = false;
3210
3211        // optimize_level_one stores the bytecode.
3212
3213        add_context.optimize_level_one = true;
3214        r.add_python_module_source_with_context(&module, &add_context)?;
3215        assert_eq!(
3216            r.resources.get(&module.name),
3217            Some(&PrePackagedResource {
3218                is_module: true,
3219                name: module.name.clone(),
3220                is_package: module.is_package,
3221                in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::FromSource(
3222                    module.source.clone()
3223                )),
3224                ..PrePackagedResource::default()
3225            })
3226        );
3227
3228        r.resources.clear();
3229        add_context.optimize_level_one = false;
3230
3231        // optimize_level_two stores the bytecode.
3232
3233        add_context.optimize_level_two = true;
3234        r.add_python_module_source_with_context(&module, &add_context)?;
3235        assert_eq!(
3236            r.resources.get(&module.name),
3237            Some(&PrePackagedResource {
3238                is_module: true,
3239                name: module.name.clone(),
3240                is_package: module.is_package,
3241                in_memory_bytecode_opt2: Some(PythonModuleBytecodeProvider::FromSource(
3242                    module.source.clone()
3243                )),
3244                ..PrePackagedResource::default()
3245            })
3246        );
3247
3248        r.resources.clear();
3249        add_context.optimize_level_two = false;
3250
3251        Ok(())
3252    }
3253
3254    #[test]
3255    fn test_add_in_memory_bytecode_module() -> Result<()> {
3256        let mut r = PythonResourceCollector::new(
3257            vec![AbstractResourceLocation::InMemory],
3258            vec![],
3259            false,
3260            false,
3261        );
3262        r.add_python_module_bytecode(
3263            &PythonModuleBytecode::new(
3264                "foo",
3265                BytecodeOptimizationLevel::Zero,
3266                false,
3267                DEFAULT_CACHE_TAG,
3268                &[42],
3269            ),
3270            &ConcreteResourceLocation::InMemory,
3271        )?;
3272
3273        assert!(r.resources.contains_key("foo"));
3274        assert_eq!(
3275            r.resources.get("foo"),
3276            Some(&PrePackagedResource {
3277                is_module: true,
3278                name: "foo".to_string(),
3279                in_memory_bytecode: Some(PythonModuleBytecodeProvider::Provided(FileData::Memory(
3280                    vec![42]
3281                ))),
3282                is_package: false,
3283                ..PrePackagedResource::default()
3284            })
3285        );
3286
3287        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3288
3289        let resources = r.compile_resources(&mut compiler)?;
3290
3291        assert_eq!(resources.resources.len(), 1);
3292        assert_eq!(
3293            resources.resources.get("foo"),
3294            Some(&Resource {
3295                is_python_module: true,
3296                name: Cow::Owned("foo".to_string()),
3297                in_memory_bytecode: Some(Cow::Owned(vec![42])),
3298                ..Resource::default()
3299            })
3300        );
3301        assert!(resources.extra_files.is_empty());
3302
3303        Ok(())
3304    }
3305
3306    #[test]
3307    fn test_add_in_memory_bytecode_module_from_source() -> Result<()> {
3308        let mut r = PythonResourceCollector::new(
3309            vec![AbstractResourceLocation::InMemory],
3310            vec![],
3311            false,
3312            false,
3313        );
3314        r.add_python_module_bytecode_from_source(
3315            &PythonModuleBytecodeFromSource {
3316                name: "foo".to_string(),
3317                source: FileData::Memory(vec![42]),
3318                optimize_level: BytecodeOptimizationLevel::Zero,
3319                is_package: false,
3320                cache_tag: DEFAULT_CACHE_TAG.to_string(),
3321                is_stdlib: false,
3322                is_test: false,
3323            },
3324            &ConcreteResourceLocation::InMemory,
3325        )?;
3326
3327        assert!(r.resources.contains_key("foo"));
3328        assert_eq!(
3329            r.resources.get("foo"),
3330            Some(&PrePackagedResource {
3331                is_module: true,
3332                name: "foo".to_string(),
3333                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
3334                    FileData::Memory(vec![42])
3335                )),
3336                is_package: false,
3337                ..PrePackagedResource::default()
3338            })
3339        );
3340
3341        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3342
3343        let resources = r.compile_resources(&mut compiler)?;
3344
3345        assert_eq!(resources.resources.len(), 1);
3346        assert_eq!(
3347            resources.resources.get("foo"),
3348            Some(&Resource {
3349                is_python_module: true,
3350                name: Cow::Owned("foo".to_string()),
3351                in_memory_bytecode: Some(Cow::Owned(b"bc0\x2a".to_vec())),
3352                ..Resource::default()
3353            })
3354        );
3355        assert!(resources.extra_files.is_empty());
3356
3357        Ok(())
3358    }
3359
3360    #[test]
3361    fn test_add_module_bytecode_with_context() -> Result<()> {
3362        let mut r = PythonResourceCollector::new(
3363            vec![AbstractResourceLocation::InMemory],
3364            vec![],
3365            false,
3366            false,
3367        );
3368
3369        let mut module = PythonModuleBytecode::new(
3370            "foo",
3371            BytecodeOptimizationLevel::Zero,
3372            false,
3373            DEFAULT_CACHE_TAG,
3374            &[42],
3375        );
3376
3377        let mut add_context = PythonResourceAddCollectionContext {
3378            include: false,
3379            location: ConcreteResourceLocation::InMemory,
3380            location_fallback: None,
3381            store_source: false,
3382            optimize_level_zero: false,
3383            optimize_level_one: false,
3384            optimize_level_two: false,
3385        };
3386
3387        // include=false is a noop.
3388        assert!(r.resources.is_empty());
3389        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3390        assert!(r.resources.is_empty());
3391
3392        add_context.include = true;
3393
3394        // optimize_level_zero=false is a noop.
3395        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3396        assert!(r.resources.is_empty());
3397
3398        // optimize_level_zero=true adds the resource.
3399        add_context.optimize_level_zero = true;
3400        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3401        assert_eq!(
3402            r.resources.get(&module.name),
3403            Some(&PrePackagedResource {
3404                is_module: true,
3405                name: module.name.clone(),
3406                is_package: module.is_package,
3407                in_memory_bytecode: Some(PythonModuleBytecodeProvider::Provided(FileData::Memory(
3408                    module.resolve_bytecode()?
3409                ))),
3410                ..PrePackagedResource::default()
3411            })
3412        );
3413
3414        r.resources.clear();
3415        add_context.optimize_level_zero = false;
3416
3417        // Other optimize_level_* fields are noop.
3418        add_context.optimize_level_one = true;
3419        add_context.optimize_level_two = true;
3420        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3421        assert!(r.resources.is_empty());
3422
3423        // No-ops for other module optimization levels.
3424        add_context.optimize_level_zero = false;
3425        add_context.optimize_level_one = false;
3426        add_context.optimize_level_two = false;
3427
3428        module.optimize_level = BytecodeOptimizationLevel::One;
3429        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3430        assert!(r.resources.is_empty());
3431        module.optimize_level = BytecodeOptimizationLevel::Two;
3432        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3433        assert!(r.resources.is_empty());
3434
3435        // optimize_level_one=true adds the resource.
3436        module.optimize_level = BytecodeOptimizationLevel::One;
3437        add_context.optimize_level_zero = true;
3438        add_context.optimize_level_one = true;
3439        add_context.optimize_level_two = true;
3440        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3441
3442        assert_eq!(
3443            r.resources.get(&module.name),
3444            Some(&PrePackagedResource {
3445                is_module: true,
3446                name: module.name.clone(),
3447                is_package: module.is_package,
3448                in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::Provided(
3449                    FileData::Memory(module.resolve_bytecode()?)
3450                )),
3451                ..PrePackagedResource::default()
3452            })
3453        );
3454
3455        r.resources.clear();
3456
3457        // optimize_level_two=true adds the resource.
3458        module.optimize_level = BytecodeOptimizationLevel::Two;
3459        r.add_python_module_bytecode_with_context(&module, &add_context)?;
3460
3461        assert_eq!(
3462            r.resources.get(&module.name),
3463            Some(&PrePackagedResource {
3464                is_module: true,
3465                name: module.name.clone(),
3466                is_package: module.is_package,
3467                in_memory_bytecode_opt2: Some(PythonModuleBytecodeProvider::Provided(
3468                    FileData::Memory(module.resolve_bytecode()?)
3469                )),
3470                ..PrePackagedResource::default()
3471            })
3472        );
3473
3474        r.resources.clear();
3475
3476        Ok(())
3477    }
3478
3479    #[test]
3480    fn test_add_in_memory_bytecode_module_parents() -> Result<()> {
3481        let mut r = PythonResourceCollector::new(
3482            vec![AbstractResourceLocation::InMemory],
3483            vec![],
3484            false,
3485            false,
3486        );
3487        r.add_python_module_bytecode_from_source(
3488            &PythonModuleBytecodeFromSource {
3489                name: "root.parent.child".to_string(),
3490                source: FileData::Memory(vec![42]),
3491                optimize_level: BytecodeOptimizationLevel::One,
3492                is_package: true,
3493                cache_tag: DEFAULT_CACHE_TAG.to_string(),
3494                is_stdlib: false,
3495                is_test: false,
3496            },
3497            &ConcreteResourceLocation::InMemory,
3498        )?;
3499
3500        assert_eq!(r.resources.len(), 1);
3501        assert_eq!(
3502            r.resources.get("root.parent.child"),
3503            Some(&PrePackagedResource {
3504                is_module: true,
3505                name: "root.parent.child".to_string(),
3506                in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::FromSource(
3507                    FileData::Memory(vec![42])
3508                )),
3509                is_package: true,
3510                ..PrePackagedResource::default()
3511            })
3512        );
3513
3514        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3515
3516        let resources = r.compile_resources(&mut compiler)?;
3517
3518        assert_eq!(resources.resources.len(), 3);
3519        assert_eq!(
3520            resources.resources.get("root"),
3521            Some(&Resource {
3522                is_python_module: true,
3523                name: Cow::Owned("root".to_string()),
3524                is_python_package: true,
3525                in_memory_bytecode_opt1: Some(Cow::Owned(b"bc1".to_vec())),
3526                ..Resource::default()
3527            })
3528        );
3529        assert_eq!(
3530            resources.resources.get("root.parent"),
3531            Some(&Resource {
3532                is_python_module: true,
3533                name: Cow::Owned("root.parent".to_string()),
3534                is_python_package: true,
3535                in_memory_bytecode_opt1: Some(Cow::Owned(b"bc1".to_vec())),
3536                ..Resource::default()
3537            })
3538        );
3539        assert_eq!(
3540            resources.resources.get("root.parent.child"),
3541            Some(&Resource {
3542                is_python_module: true,
3543                name: Cow::Owned("root.parent.child".to_string()),
3544                is_python_package: true,
3545                in_memory_bytecode_opt1: Some(Cow::Owned(b"bc1\x2a".to_vec())),
3546                ..Resource::default()
3547            })
3548        );
3549        assert!(resources.extra_files.is_empty());
3550
3551        Ok(())
3552    }
3553
3554    #[test]
3555    fn test_add_module_bytecode_from_source_with_context() -> Result<()> {
3556        let mut r = PythonResourceCollector::new(
3557            vec![AbstractResourceLocation::InMemory],
3558            vec![],
3559            false,
3560            false,
3561        );
3562
3563        let mut module = PythonModuleBytecodeFromSource {
3564            name: "foo".to_string(),
3565            source: FileData::Memory(vec![42]),
3566            optimize_level: BytecodeOptimizationLevel::Zero,
3567            is_package: false,
3568            cache_tag: DEFAULT_CACHE_TAG.to_string(),
3569            is_stdlib: false,
3570            is_test: false,
3571        };
3572
3573        let mut add_context = PythonResourceAddCollectionContext {
3574            include: false,
3575            location: ConcreteResourceLocation::InMemory,
3576            location_fallback: None,
3577            store_source: false,
3578            optimize_level_zero: false,
3579            optimize_level_one: false,
3580            optimize_level_two: false,
3581        };
3582
3583        // include=false is a noop.
3584        assert!(r.resources.is_empty());
3585        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3586        assert!(r.resources.is_empty());
3587
3588        add_context.include = true;
3589
3590        // optimize_level_zero=false is a noop.
3591        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3592        assert!(r.resources.is_empty());
3593
3594        // optimize_level_zero=true adds the resource.
3595        add_context.optimize_level_zero = true;
3596        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3597        assert_eq!(
3598            r.resources.get(&module.name),
3599            Some(&PrePackagedResource {
3600                is_module: true,
3601                name: module.name.clone(),
3602                is_package: module.is_package,
3603                in_memory_bytecode: Some(PythonModuleBytecodeProvider::FromSource(
3604                    module.source.clone()
3605                )),
3606                ..PrePackagedResource::default()
3607            })
3608        );
3609
3610        r.resources.clear();
3611        add_context.optimize_level_zero = false;
3612
3613        // Other optimize_level_* fields are noop.
3614        add_context.optimize_level_one = true;
3615        add_context.optimize_level_two = true;
3616        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3617        assert!(r.resources.is_empty());
3618
3619        // No-ops for other module optimization levels.
3620        add_context.optimize_level_zero = false;
3621        add_context.optimize_level_one = false;
3622        add_context.optimize_level_two = false;
3623
3624        module.optimize_level = BytecodeOptimizationLevel::One;
3625        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3626        assert!(r.resources.is_empty());
3627        module.optimize_level = BytecodeOptimizationLevel::Two;
3628        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3629        assert!(r.resources.is_empty());
3630
3631        // optimize_level_one=true adds the resource.
3632        module.optimize_level = BytecodeOptimizationLevel::One;
3633        add_context.optimize_level_zero = true;
3634        add_context.optimize_level_one = true;
3635        add_context.optimize_level_two = true;
3636        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3637
3638        assert_eq!(
3639            r.resources.get(&module.name),
3640            Some(&PrePackagedResource {
3641                is_module: true,
3642                name: module.name.clone(),
3643                is_package: module.is_package,
3644                in_memory_bytecode_opt1: Some(PythonModuleBytecodeProvider::FromSource(
3645                    module.source.clone()
3646                )),
3647                ..PrePackagedResource::default()
3648            })
3649        );
3650
3651        r.resources.clear();
3652
3653        // optimize_level_two=true adds the resource.
3654        module.optimize_level = BytecodeOptimizationLevel::Two;
3655        r.add_python_module_bytecode_from_source_with_context(&module, &add_context)?;
3656
3657        assert_eq!(
3658            r.resources.get(&module.name),
3659            Some(&PrePackagedResource {
3660                is_module: true,
3661                name: module.name.clone(),
3662                is_package: module.is_package,
3663                in_memory_bytecode_opt2: Some(PythonModuleBytecodeProvider::FromSource(
3664                    module.source.clone()
3665                )),
3666                ..PrePackagedResource::default()
3667            })
3668        );
3669
3670        r.resources.clear();
3671
3672        Ok(())
3673    }
3674
3675    #[test]
3676    fn test_add_in_memory_package_resource() -> Result<()> {
3677        let mut r = PythonResourceCollector::new(
3678            vec![AbstractResourceLocation::InMemory],
3679            vec![],
3680            false,
3681            false,
3682        );
3683        r.add_python_package_resource(
3684            &PythonPackageResource {
3685                leaf_package: "foo".to_string(),
3686                relative_name: "resource.txt".to_string(),
3687                data: FileData::Memory(vec![42]),
3688                is_stdlib: false,
3689                is_test: false,
3690            },
3691            &ConcreteResourceLocation::InMemory,
3692        )?;
3693
3694        assert_eq!(r.resources.len(), 1);
3695        assert_eq!(
3696            r.resources.get("foo"),
3697            Some(&PrePackagedResource {
3698                is_module: true,
3699                name: "foo".to_string(),
3700                is_package: true,
3701                in_memory_resources: Some(
3702                    [("resource.txt".to_string(), FileData::Memory(vec![42]))]
3703                        .iter()
3704                        .cloned()
3705                        .collect()
3706                ),
3707                ..PrePackagedResource::default()
3708            })
3709        );
3710
3711        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3712
3713        let resources = r.compile_resources(&mut compiler)?;
3714
3715        assert_eq!(resources.resources.len(), 1);
3716        assert_eq!(
3717            resources.resources.get("foo"),
3718            Some(&Resource {
3719                is_python_module: true,
3720                name: Cow::Owned("foo".to_string()),
3721                is_python_package: true,
3722                in_memory_package_resources: Some(
3723                    [(Cow::Owned("resource.txt".to_string()), Cow::Owned(vec![42]))]
3724                        .iter()
3725                        .cloned()
3726                        .collect()
3727                ),
3728                ..Resource::default()
3729            })
3730        );
3731        assert!(resources.extra_files.is_empty());
3732
3733        Ok(())
3734    }
3735
3736    #[test]
3737    fn test_add_relative_path_package_resource() -> Result<()> {
3738        let mut r = PythonResourceCollector::new(
3739            vec![AbstractResourceLocation::RelativePath],
3740            vec![],
3741            false,
3742            false,
3743        );
3744        r.add_python_package_resource(
3745            &PythonPackageResource {
3746                leaf_package: "foo".to_string(),
3747                relative_name: "resource.txt".to_string(),
3748                data: FileData::Memory(vec![42]),
3749                is_stdlib: false,
3750                is_test: false,
3751            },
3752            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
3753        )?;
3754
3755        assert_eq!(r.resources.len(), 1);
3756        assert_eq!(
3757            r.resources.get("foo"),
3758            Some(&PrePackagedResource {
3759                is_module: true,
3760                name: "foo".to_string(),
3761                is_package: true,
3762                relative_path_package_resources: Some(
3763                    [(
3764                        "resource.txt".to_string(),
3765                        (
3766                            PathBuf::from("prefix/foo/resource.txt"),
3767                            FileData::Memory(vec![42])
3768                        )
3769                    )]
3770                    .iter()
3771                    .cloned()
3772                    .collect()
3773                ),
3774                ..PrePackagedResource::default()
3775            })
3776        );
3777
3778        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3779
3780        let resources = r.compile_resources(&mut compiler)?;
3781
3782        assert_eq!(resources.resources.len(), 1);
3783        assert_eq!(
3784            resources.resources.get("foo"),
3785            Some(&Resource {
3786                is_python_module: true,
3787                name: Cow::Owned("foo".to_string()),
3788                is_python_package: true,
3789                relative_path_package_resources: Some(
3790                    [(
3791                        Cow::Owned("resource.txt".to_string()),
3792                        Cow::Owned(PathBuf::from("prefix/foo/resource.txt")),
3793                    )]
3794                    .iter()
3795                    .cloned()
3796                    .collect()
3797                ),
3798                ..Resource::default()
3799            })
3800        );
3801        assert_eq!(
3802            resources.extra_files,
3803            vec![(
3804                PathBuf::from("prefix/foo/resource.txt"),
3805                FileData::Memory(vec![42]),
3806                false
3807            ),]
3808        );
3809
3810        Ok(())
3811    }
3812
3813    #[test]
3814    fn test_add_package_resource_with_context() -> Result<()> {
3815        let mut r = PythonResourceCollector::new(
3816            vec![AbstractResourceLocation::InMemory],
3817            vec![],
3818            false,
3819            false,
3820        );
3821
3822        let resource = PythonPackageResource {
3823            leaf_package: "foo".to_string(),
3824            relative_name: "bar.txt".to_string(),
3825            data: FileData::Memory(vec![42]),
3826            is_stdlib: false,
3827            is_test: false,
3828        };
3829
3830        let mut add_context = PythonResourceAddCollectionContext {
3831            include: false,
3832            location: ConcreteResourceLocation::InMemory,
3833            location_fallback: None,
3834            store_source: false,
3835            optimize_level_zero: false,
3836            optimize_level_one: false,
3837            optimize_level_two: false,
3838        };
3839
3840        // include=false is a noop.
3841        assert!(r.resources.is_empty());
3842        r.add_python_package_resource_with_context(&resource, &add_context)?;
3843        assert!(r.resources.is_empty());
3844
3845        // include=true adds the resource.
3846        add_context.include = true;
3847        r.add_python_package_resource_with_context(&resource, &add_context)?;
3848        assert_eq!(
3849            r.resources.get(&resource.leaf_package),
3850            Some(&PrePackagedResource {
3851                is_module: true,
3852                name: resource.leaf_package.clone(),
3853                is_package: true,
3854                in_memory_resources: Some(
3855                    [(resource.relative_name.clone(), resource.data.clone())]
3856                        .iter()
3857                        .cloned()
3858                        .collect()
3859                ),
3860                ..PrePackagedResource::default()
3861            })
3862        );
3863
3864        r.resources.clear();
3865
3866        // location_fallback works.
3867        r.allowed_locations = vec![AbstractResourceLocation::RelativePath];
3868        add_context.location_fallback =
3869            Some(ConcreteResourceLocation::RelativePath("prefix".to_string()));
3870        r.add_python_package_resource_with_context(&resource, &add_context)?;
3871        assert_eq!(
3872            r.resources.get(&resource.leaf_package),
3873            Some(&PrePackagedResource {
3874                is_module: true,
3875                name: resource.leaf_package.clone(),
3876                is_package: true,
3877                relative_path_package_resources: Some(
3878                    [(
3879                        resource.relative_name.clone(),
3880                        (
3881                            PathBuf::from("prefix")
3882                                .join(resource.leaf_package)
3883                                .join(resource.relative_name),
3884                            resource.data.clone()
3885                        )
3886                    )]
3887                    .iter()
3888                    .cloned()
3889                    .collect()
3890                ),
3891                ..PrePackagedResource::default()
3892            })
3893        );
3894
3895        r.resources.clear();
3896
3897        Ok(())
3898    }
3899
3900    #[test]
3901    fn test_add_in_memory_package_distribution_resource() -> Result<()> {
3902        let mut r = PythonResourceCollector::new(
3903            vec![AbstractResourceLocation::InMemory],
3904            vec![],
3905            false,
3906            false,
3907        );
3908        r.add_python_package_distribution_resource(
3909            &PythonPackageDistributionResource {
3910                location: PythonPackageDistributionResourceFlavor::DistInfo,
3911                package: "mypackage".to_string(),
3912                version: "1.0".to_string(),
3913                name: "resource.txt".to_string(),
3914                data: FileData::Memory(vec![42]),
3915            },
3916            &ConcreteResourceLocation::InMemory,
3917        )?;
3918
3919        assert_eq!(r.resources.len(), 1);
3920        assert_eq!(
3921            r.resources.get("mypackage"),
3922            Some(&PrePackagedResource {
3923                is_module: true,
3924                name: "mypackage".to_string(),
3925                is_package: true,
3926                in_memory_distribution_resources: Some(
3927                    [("resource.txt".to_string(), FileData::Memory(vec![42]))]
3928                        .iter()
3929                        .cloned()
3930                        .collect()
3931                ),
3932                ..PrePackagedResource::default()
3933            })
3934        );
3935
3936        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
3937
3938        let resources = r.compile_resources(&mut compiler)?;
3939
3940        assert_eq!(resources.resources.len(), 1);
3941        assert_eq!(
3942            resources.resources.get("mypackage"),
3943            Some(&Resource {
3944                is_python_module: true,
3945                name: Cow::Owned("mypackage".to_string()),
3946                is_python_package: true,
3947                in_memory_distribution_resources: Some(
3948                    [(Cow::Owned("resource.txt".to_string()), Cow::Owned(vec![42]))]
3949                        .iter()
3950                        .cloned()
3951                        .collect()
3952                ),
3953                ..Resource::default()
3954            })
3955        );
3956        assert!(resources.extra_files.is_empty());
3957
3958        Ok(())
3959    }
3960
3961    #[test]
3962    fn test_add_relative_path_package_distribution_resource() -> Result<()> {
3963        let mut r = PythonResourceCollector::new(
3964            vec![AbstractResourceLocation::RelativePath],
3965            vec![],
3966            false,
3967            false,
3968        );
3969        r.add_python_package_distribution_resource(
3970            &PythonPackageDistributionResource {
3971                location: PythonPackageDistributionResourceFlavor::DistInfo,
3972                package: "mypackage".to_string(),
3973                version: "1.0".to_string(),
3974                name: "resource.txt".to_string(),
3975                data: FileData::Memory(vec![42]),
3976            },
3977            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
3978        )?;
3979
3980        assert_eq!(r.resources.len(), 1);
3981        assert_eq!(
3982            r.resources.get("mypackage"),
3983            Some(&PrePackagedResource {
3984                is_module: true,
3985                name: "mypackage".to_string(),
3986                is_package: true,
3987                relative_path_distribution_resources: Some(
3988                    [(
3989                        "resource.txt".to_string(),
3990                        (
3991                            PathBuf::from("prefix/mypackage-1.0.dist-info/resource.txt"),
3992                            FileData::Memory(vec![42])
3993                        )
3994                    )]
3995                    .iter()
3996                    .cloned()
3997                    .collect()
3998                ),
3999                ..PrePackagedResource::default()
4000            })
4001        );
4002
4003        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4004
4005        let resources = r.compile_resources(&mut compiler)?;
4006
4007        assert_eq!(resources.resources.len(), 1);
4008        assert_eq!(
4009            resources.resources.get("mypackage"),
4010            Some(&Resource {
4011                is_python_module: true,
4012                name: Cow::Owned("mypackage".to_string()),
4013                is_python_package: true,
4014                relative_path_distribution_resources: Some(
4015                    [(
4016                        Cow::Owned("resource.txt".to_string()),
4017                        Cow::Owned(PathBuf::from("prefix/mypackage-1.0.dist-info/resource.txt")),
4018                    )]
4019                    .iter()
4020                    .cloned()
4021                    .collect()
4022                ),
4023                ..Resource::default()
4024            })
4025        );
4026        assert_eq!(
4027            resources.extra_files,
4028            vec![(
4029                PathBuf::from("prefix/mypackage-1.0.dist-info/resource.txt"),
4030                FileData::Memory(vec![42]),
4031                false
4032            ),]
4033        );
4034
4035        Ok(())
4036    }
4037
4038    #[test]
4039    fn test_add_package_distribution_resource_with_context() -> Result<()> {
4040        let mut r = PythonResourceCollector::new(
4041            vec![AbstractResourceLocation::InMemory],
4042            vec![],
4043            false,
4044            false,
4045        );
4046
4047        let resource = PythonPackageDistributionResource {
4048            location: PythonPackageDistributionResourceFlavor::DistInfo,
4049            package: "foo".to_string(),
4050            version: "1.0".to_string(),
4051            name: "resource.txt".to_string(),
4052            data: FileData::Memory(vec![42]),
4053        };
4054
4055        let mut add_context = PythonResourceAddCollectionContext {
4056            include: false,
4057            location: ConcreteResourceLocation::InMemory,
4058            location_fallback: None,
4059            store_source: false,
4060            optimize_level_zero: false,
4061            optimize_level_one: false,
4062            optimize_level_two: false,
4063        };
4064
4065        // include=false is a noop.
4066        assert!(r.resources.is_empty());
4067        r.add_python_package_distribution_resource_with_context(&resource, &add_context)?;
4068        assert!(r.resources.is_empty());
4069
4070        // include=true adds the resource.
4071        add_context.include = true;
4072        r.add_python_package_distribution_resource_with_context(&resource, &add_context)?;
4073        assert_eq!(
4074            r.resources.get(&resource.package),
4075            Some(&PrePackagedResource {
4076                is_module: true,
4077                name: resource.package.clone(),
4078                is_package: true,
4079                in_memory_distribution_resources: Some(
4080                    [(resource.name.clone(), resource.data.clone())]
4081                        .iter()
4082                        .cloned()
4083                        .collect()
4084                ),
4085                ..PrePackagedResource::default()
4086            })
4087        );
4088
4089        r.resources.clear();
4090
4091        // location_fallback works.
4092        r.allowed_locations = vec![AbstractResourceLocation::RelativePath];
4093        add_context.location_fallback =
4094            Some(ConcreteResourceLocation::RelativePath("prefix".to_string()));
4095        r.add_python_package_distribution_resource_with_context(&resource, &add_context)?;
4096        assert_eq!(
4097            r.resources.get(&resource.package),
4098            Some(&PrePackagedResource {
4099                is_module: true,
4100                name: resource.package.clone(),
4101                is_package: true,
4102                relative_path_distribution_resources: Some(
4103                    [(
4104                        resource.name.clone(),
4105                        (resource.resolve_path("prefix"), resource.data.clone())
4106                    )]
4107                    .iter()
4108                    .cloned()
4109                    .collect()
4110                ),
4111                ..PrePackagedResource::default()
4112            })
4113        );
4114
4115        r.resources.clear();
4116
4117        Ok(())
4118    }
4119
4120    #[test]
4121    fn test_add_builtin_python_extension_module() -> Result<()> {
4122        let mut c = PythonResourceCollector::new(
4123            vec![AbstractResourceLocation::InMemory],
4124            vec![AbstractResourceLocation::InMemory],
4125            false,
4126            false,
4127        );
4128
4129        let em = PythonExtensionModule {
4130            name: "_io".to_string(),
4131            init_fn: Some("PyInit__io".to_string()),
4132            extension_file_suffix: "".to_string(),
4133            shared_library: None,
4134            object_file_data: vec![],
4135            is_package: false,
4136            link_libraries: vec![],
4137            is_stdlib: true,
4138            builtin_default: true,
4139            required: true,
4140            variant: None,
4141            license: None,
4142        };
4143
4144        c.add_builtin_python_extension_module(&em)?;
4145        assert_eq!(c.resources.len(), 1);
4146        assert_eq!(
4147            c.resources.get("_io"),
4148            Some(&PrePackagedResource {
4149                is_builtin_extension_module: true,
4150                name: "_io".to_string(),
4151                ..PrePackagedResource::default()
4152            })
4153        );
4154
4155        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4156
4157        let resources = c.compile_resources(&mut compiler)?;
4158
4159        assert_eq!(resources.resources.len(), 1);
4160        assert_eq!(
4161            resources.resources.get("_io"),
4162            Some(&Resource {
4163                is_python_builtin_extension_module: true,
4164                name: Cow::Owned("_io".to_string()),
4165                ..Resource::default()
4166            })
4167        );
4168        assert!(resources.extra_files.is_empty());
4169
4170        Ok(())
4171    }
4172
4173    #[test]
4174    fn test_add_in_memory_python_extension_module_shared_library() -> Result<()> {
4175        let em = PythonExtensionModule {
4176            name: "myext".to_string(),
4177            init_fn: Some("PyInit__myext".to_string()),
4178            extension_file_suffix: ".so".to_string(),
4179            shared_library: Some(FileData::Memory(vec![42])),
4180            object_file_data: vec![],
4181            is_package: false,
4182            link_libraries: vec![LibraryDependency {
4183                name: "foo".to_string(),
4184                static_library: None,
4185                static_filename: None,
4186                dynamic_library: Some(FileData::Memory(vec![40])),
4187                dynamic_filename: Some(PathBuf::from("libfoo.so")),
4188                framework: false,
4189                system: false,
4190            }],
4191            is_stdlib: false,
4192            builtin_default: false,
4193            required: false,
4194            variant: None,
4195            license: None,
4196        };
4197
4198        let mut c = PythonResourceCollector::new(
4199            vec![AbstractResourceLocation::InMemory],
4200            vec![],
4201            false,
4202            false,
4203        );
4204
4205        let res = c.add_python_extension_module(&em, &ConcreteResourceLocation::InMemory);
4206        assert!(res.is_err());
4207        assert_eq!(res.err().unwrap().to_string(), "cannot add extension module myext for in-memory import because in-memory loading is not supported/allowed");
4208
4209        let mut c = PythonResourceCollector::new(
4210            vec![AbstractResourceLocation::InMemory],
4211            vec![AbstractResourceLocation::InMemory],
4212            false,
4213            false,
4214        );
4215
4216        c.add_python_extension_module(&em, &ConcreteResourceLocation::InMemory)?;
4217        assert_eq!(c.resources.len(), 2);
4218        assert_eq!(
4219            c.resources.get("myext"),
4220            Some(&PrePackagedResource {
4221                is_extension_module: true,
4222                name: "myext".to_string(),
4223                in_memory_extension_module_shared_library: Some(FileData::Memory(vec![42])),
4224                shared_library_dependency_names: Some(vec!["foo".to_string()]),
4225                ..PrePackagedResource::default()
4226            })
4227        );
4228        assert_eq!(
4229            c.resources.get("foo"),
4230            Some(&PrePackagedResource {
4231                is_shared_library: true,
4232                name: "foo".to_string(),
4233                in_memory_shared_library: Some(FileData::Memory(vec![40])),
4234                ..PrePackagedResource::default()
4235            })
4236        );
4237
4238        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4239
4240        let resources = c.compile_resources(&mut compiler)?;
4241
4242        assert_eq!(resources.resources.len(), 2);
4243        assert_eq!(
4244            resources.resources.get("myext"),
4245            Some(&Resource {
4246                is_python_extension_module: true,
4247                name: Cow::Owned("myext".to_string()),
4248                in_memory_extension_module_shared_library: Some(Cow::Owned(vec![42])),
4249                shared_library_dependency_names: Some(vec![Cow::Owned("foo".to_string())]),
4250                ..Resource::default()
4251            })
4252        );
4253        assert_eq!(
4254            resources.resources.get("foo"),
4255            Some(&Resource {
4256                is_shared_library: true,
4257                name: Cow::Owned("foo".to_string()),
4258                in_memory_shared_library: Some(Cow::Owned(vec![40])),
4259                ..Resource::default()
4260            })
4261        );
4262
4263        assert!(resources.extra_files.is_empty());
4264
4265        Ok(())
4266    }
4267
4268    #[test]
4269    fn test_add_relative_path_python_extension_module() -> Result<()> {
4270        let em = PythonExtensionModule {
4271            name: "foo.bar".to_string(),
4272            init_fn: None,
4273            extension_file_suffix: ".so".to_string(),
4274            shared_library: Some(FileData::Memory(vec![42])),
4275            object_file_data: vec![],
4276            is_package: false,
4277            link_libraries: vec![LibraryDependency {
4278                name: "mylib".to_string(),
4279                static_library: None,
4280                static_filename: None,
4281                dynamic_library: Some(FileData::Memory(vec![40])),
4282                dynamic_filename: Some(PathBuf::from("libmylib.so")),
4283                framework: false,
4284                system: false,
4285            }],
4286            is_stdlib: false,
4287            builtin_default: false,
4288            required: false,
4289            variant: None,
4290            license: None,
4291        };
4292
4293        let mut c = PythonResourceCollector::new(
4294            vec![AbstractResourceLocation::RelativePath],
4295            vec![],
4296            false,
4297            false,
4298        );
4299        let res = c.add_python_extension_module(
4300            &em,
4301            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
4302        );
4303        assert!(res.is_err());
4304        assert_eq!(res.err().unwrap().to_string(), "cannot add extension module foo.bar as a file because extension modules as files are not allowed");
4305
4306        let mut c = PythonResourceCollector::new(
4307            vec![AbstractResourceLocation::RelativePath],
4308            vec![AbstractResourceLocation::RelativePath],
4309            false,
4310            false,
4311        );
4312
4313        c.add_python_extension_module(
4314            &em,
4315            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
4316        )?;
4317        assert_eq!(c.resources.len(), 2);
4318        assert_eq!(
4319            c.resources.get("foo.bar"),
4320            Some(&PrePackagedResource {
4321                is_extension_module: true,
4322                name: "foo.bar".to_string(),
4323                is_package: false,
4324                relative_path_extension_module_shared_library: Some((
4325                    PathBuf::from("prefix/foo/bar.so"),
4326                    FileData::Memory(vec![42])
4327                )),
4328                shared_library_dependency_names: Some(vec!["mylib".to_string()]),
4329                ..PrePackagedResource::default()
4330            })
4331        );
4332        assert_eq!(
4333            c.resources.get("mylib"),
4334            Some(&PrePackagedResource {
4335                is_shared_library: true,
4336                name: "mylib".to_string(),
4337                relative_path_shared_library: Some((
4338                    "prefix/foo".to_string(),
4339                    PathBuf::from("libmylib.so"),
4340                    FileData::Memory(vec![40])
4341                )),
4342                ..PrePackagedResource::default()
4343            })
4344        );
4345
4346        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4347
4348        let resources = c.compile_resources(&mut compiler)?;
4349
4350        assert_eq!(resources.resources.len(), 3);
4351        assert_eq!(
4352            resources.resources.get("foo"),
4353            Some(&Resource {
4354                is_python_module: true,
4355                name: Cow::Owned("foo".to_string()),
4356                is_python_package: true,
4357                ..Resource::default()
4358            })
4359        );
4360        assert_eq!(
4361            resources.resources.get("foo.bar"),
4362            Some(&Resource {
4363                is_python_extension_module: true,
4364                name: Cow::Owned("foo.bar".to_string()),
4365                is_python_package: false,
4366                relative_path_extension_module_shared_library: Some(Cow::Owned(PathBuf::from(
4367                    "prefix/foo/bar.so"
4368                ))),
4369                shared_library_dependency_names: Some(vec![Cow::Owned("mylib".to_string())]),
4370                ..Resource::default()
4371            })
4372        );
4373        assert_eq!(
4374            resources.resources.get("mylib"),
4375            Some(&Resource {
4376                is_shared_library: true,
4377                name: Cow::Owned("mylib".to_string()),
4378                ..Resource::default()
4379            })
4380        );
4381
4382        assert_eq!(
4383            resources.extra_files,
4384            vec![
4385                (
4386                    PathBuf::from("prefix/foo/bar.so"),
4387                    FileData::Memory(vec![42]),
4388                    true
4389                ),
4390                (
4391                    PathBuf::from("prefix/foo/libmylib.so"),
4392                    FileData::Memory(vec![40]),
4393                    true
4394                )
4395            ]
4396        );
4397
4398        Ok(())
4399    }
4400
4401    #[test]
4402    fn test_add_shared_library_and_module() -> Result<()> {
4403        let mut r = PythonResourceCollector::new(
4404            vec![AbstractResourceLocation::InMemory],
4405            vec![],
4406            false,
4407            false,
4408        );
4409
4410        r.add_python_module_source(
4411            &PythonModuleSource {
4412                name: "foo".to_string(),
4413                source: FileData::Memory(vec![1]),
4414                is_package: true,
4415                cache_tag: DEFAULT_CACHE_TAG.to_string(),
4416                is_stdlib: false,
4417                is_test: false,
4418            },
4419            &ConcreteResourceLocation::InMemory,
4420        )?;
4421
4422        r.add_shared_library(
4423            &SharedLibrary {
4424                name: "foo".to_string(),
4425                data: FileData::Memory(vec![2]),
4426                filename: None,
4427            },
4428            &ConcreteResourceLocation::InMemory,
4429        )?;
4430
4431        assert_eq!(r.resources.len(), 1);
4432        assert_eq!(
4433            r.resources.get("foo"),
4434            Some(&PrePackagedResource {
4435                is_module: true,
4436                is_shared_library: true,
4437                name: "foo".to_string(),
4438                is_package: true,
4439                in_memory_source: Some(FileData::Memory(vec![1])),
4440                in_memory_shared_library: Some(FileData::Memory(vec![2])),
4441                ..PrePackagedResource::default()
4442            })
4443        );
4444
4445        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4446
4447        let resources = r.compile_resources(&mut compiler)?;
4448
4449        assert_eq!(resources.resources.len(), 1);
4450        assert_eq!(
4451            resources.resources.get("foo"),
4452            Some(&Resource {
4453                is_python_module: true,
4454                is_shared_library: true,
4455                name: Cow::Owned("foo".to_string()),
4456                is_python_package: true,
4457                in_memory_source: Some(Cow::Owned(vec![1])),
4458                in_memory_shared_library: Some(Cow::Owned(vec![2])),
4459                ..Resource::default()
4460            })
4461        );
4462
4463        Ok(())
4464    }
4465
4466    #[test]
4467    fn test_add_in_memory_file_data() -> Result<()> {
4468        let mut r = PythonResourceCollector::new(
4469            vec![AbstractResourceLocation::InMemory],
4470            vec![],
4471            false,
4472            false,
4473        );
4474        assert!(r
4475            .add_file_data(
4476                &File::new("foo/bar.py", vec![42]),
4477                &ConcreteResourceLocation::InMemory,
4478            )
4479            .is_err());
4480
4481        r.allow_files = true;
4482        r.add_file_data(
4483            &File::new("foo/bar.py", vec![42]),
4484            &ConcreteResourceLocation::InMemory,
4485        )?;
4486
4487        assert!(r.resources.contains_key("foo/bar.py"));
4488        assert_eq!(
4489            r.resources.get("foo/bar.py"),
4490            Some(&PrePackagedResource {
4491                is_utf8_filename_data: true,
4492                name: "foo/bar.py".to_string(),
4493                file_data_embedded: Some(FileData::Memory(vec![42])),
4494                ..PrePackagedResource::default()
4495            })
4496        );
4497
4498        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4499
4500        let resources = r.compile_resources(&mut compiler)?;
4501
4502        assert_eq!(resources.resources.len(), 1);
4503        assert_eq!(
4504            resources.resources.get("foo/bar.py"),
4505            Some(&Resource {
4506                is_utf8_filename_data: true,
4507                name: Cow::Owned("foo/bar.py".to_string()),
4508                file_data_embedded: Some(Cow::Owned(vec![42])),
4509                ..Resource::default()
4510            })
4511        );
4512        assert!(resources.extra_files.is_empty());
4513
4514        Ok(())
4515    }
4516
4517    #[test]
4518    fn test_add_relative_path_file_data() -> Result<()> {
4519        let mut r = PythonResourceCollector::new(
4520            vec![AbstractResourceLocation::RelativePath],
4521            vec![],
4522            false,
4523            true,
4524        );
4525        r.add_file_data(
4526            &File::new("foo/bar.py", vec![42]),
4527            &ConcreteResourceLocation::RelativePath("prefix".to_string()),
4528        )?;
4529
4530        assert!(r.resources.contains_key("foo/bar.py"));
4531        assert_eq!(
4532            r.resources.get("foo/bar.py"),
4533            Some(&PrePackagedResource {
4534                is_utf8_filename_data: true,
4535                name: "foo/bar.py".to_string(),
4536                file_data_utf8_relative_path: Some((
4537                    PathBuf::from("prefix/foo/bar.py"),
4538                    FileData::Memory(vec![42])
4539                )),
4540                ..PrePackagedResource::default()
4541            })
4542        );
4543
4544        let mut compiler = FakeBytecodeCompiler { magic_number: 42 };
4545
4546        let resources = r.compile_resources(&mut compiler)?;
4547
4548        assert_eq!(resources.resources.len(), 1);
4549        assert_eq!(
4550            resources.resources.get("foo/bar.py"),
4551            Some(&Resource {
4552                is_utf8_filename_data: true,
4553                name: Cow::Owned("foo/bar.py".to_string()),
4554                file_data_utf8_relative_path: Some(Cow::Owned("prefix/foo/bar.py".to_string())),
4555                ..Resource::default()
4556            })
4557        );
4558        assert_eq!(
4559            resources.extra_files,
4560            vec![(
4561                PathBuf::from("prefix/foo/bar.py"),
4562                FileData::Memory(vec![42]),
4563                false
4564            )]
4565        );
4566
4567        Ok(())
4568    }
4569
4570    #[test]
4571    fn test_add_file_data_with_context() -> Result<()> {
4572        let mut r = PythonResourceCollector::new(
4573            vec![AbstractResourceLocation::InMemory],
4574            vec![],
4575            false,
4576            true,
4577        );
4578
4579        let file = File::new("foo/bar.py", FileEntry::new_from_data(vec![42], true));
4580
4581        let mut add_context = PythonResourceAddCollectionContext {
4582            include: false,
4583            location: ConcreteResourceLocation::InMemory,
4584            location_fallback: None,
4585            store_source: false,
4586            optimize_level_zero: false,
4587            optimize_level_one: false,
4588            optimize_level_two: false,
4589        };
4590
4591        // include=false is a noop.
4592        assert!(r.resources.is_empty());
4593        r.add_file_data_with_context(&file, &add_context)?;
4594        assert!(r.resources.is_empty());
4595
4596        // include=true adds the resource.
4597        add_context.include = true;
4598        r.add_file_data_with_context(&file, &add_context)?;
4599        assert_eq!(
4600            r.resources.get(&file.path_string()),
4601            Some(&PrePackagedResource {
4602                name: file.path_string(),
4603                is_utf8_filename_data: true,
4604                file_executable: true,
4605                file_data_embedded: Some(file.entry().file_data().clone()),
4606                ..PrePackagedResource::default()
4607            })
4608        );
4609        r.resources.clear();
4610
4611        // location_fallback works.
4612        r.allowed_locations = vec![AbstractResourceLocation::RelativePath];
4613        add_context.location_fallback =
4614            Some(ConcreteResourceLocation::RelativePath("prefix".to_string()));
4615        r.add_file_data_with_context(&file, &add_context)?;
4616        assert_eq!(
4617            r.resources.get(&file.path_string()),
4618            Some(&PrePackagedResource {
4619                name: file.path_string(),
4620                is_utf8_filename_data: true,
4621                file_executable: true,
4622                file_data_utf8_relative_path: Some((
4623                    PathBuf::from("prefix").join(file.path_string()),
4624                    file.entry().file_data().clone()
4625                )),
4626                ..PrePackagedResource::default()
4627            })
4628        );
4629
4630        Ok(())
4631    }
4632
4633    #[test]
4634    fn test_find_dunder_file() -> Result<()> {
4635        let mut r = PythonResourceCollector::new(
4636            vec![AbstractResourceLocation::InMemory],
4637            vec![],
4638            false,
4639            false,
4640        );
4641        assert_eq!(r.find_dunder_file()?.len(), 0);
4642
4643        r.add_python_module_source(
4644            &PythonModuleSource {
4645                name: "foo.bar".to_string(),
4646                source: FileData::Memory(vec![]),
4647                is_package: false,
4648                cache_tag: DEFAULT_CACHE_TAG.to_string(),
4649                is_stdlib: false,
4650                is_test: false,
4651            },
4652            &ConcreteResourceLocation::InMemory,
4653        )?;
4654        assert_eq!(r.find_dunder_file()?.len(), 0);
4655
4656        r.add_python_module_source(
4657            &PythonModuleSource {
4658                name: "baz".to_string(),
4659                source: FileData::Memory(Vec::from("import foo; if __file__ == 'ignored'")),
4660                is_package: false,
4661                cache_tag: DEFAULT_CACHE_TAG.to_string(),
4662                is_stdlib: false,
4663                is_test: false,
4664            },
4665            &ConcreteResourceLocation::InMemory,
4666        )?;
4667        assert_eq!(r.find_dunder_file()?.len(), 1);
4668        assert!(r.find_dunder_file()?.contains("baz"));
4669
4670        r.add_python_module_bytecode_from_source(
4671            &PythonModuleBytecodeFromSource {
4672                name: "bytecode".to_string(),
4673                source: FileData::Memory(Vec::from("import foo; if __file__")),
4674                optimize_level: BytecodeOptimizationLevel::Zero,
4675                is_package: false,
4676                cache_tag: DEFAULT_CACHE_TAG.to_string(),
4677                is_stdlib: false,
4678                is_test: false,
4679            },
4680            &ConcreteResourceLocation::InMemory,
4681        )?;
4682        assert_eq!(r.find_dunder_file()?.len(), 2);
4683        assert!(r.find_dunder_file()?.contains("bytecode"));
4684
4685        Ok(())
4686    }
4687}