python_packaging/
module_util.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/*! Utility functions related to Python modules. */
10
11use std::{collections::BTreeSet, path::Path, path::PathBuf};
12
13/// Represents file name suffixes for Python modules.
14#[derive(Clone, Debug, PartialEq, Eq)]
15pub struct PythonModuleSuffixes {
16    /// Suffixes for Python source modules.
17    pub source: Vec<String>,
18
19    /// Suffixes for Python bytecode modules.
20    pub bytecode: Vec<String>,
21
22    /// Suffixes for Python debug bytecode modules.
23    pub debug_bytecode: Vec<String>,
24
25    /// Suffixes for Python optimized bytecode modules.
26    pub optimized_bytecode: Vec<String>,
27
28    /// Suffixes for Python extension modules.
29    pub extension: Vec<String>,
30}
31
32/// Resolve the set of packages present in a fully qualified module name.
33pub fn packages_from_module_name(module: &str) -> BTreeSet<String> {
34    let mut package_names = BTreeSet::new();
35
36    let mut search: &str = module;
37
38    while let Some(idx) = search.rfind('.') {
39        package_names.insert(search[0..idx].to_string());
40        search = &search[0..idx];
41    }
42
43    package_names
44}
45
46/// Resolve the set of packages present in a series of fully qualified module names.
47pub fn packages_from_module_names<I>(names: I) -> BTreeSet<String>
48where
49    I: Iterator<Item = String>,
50{
51    let mut package_names = BTreeSet::new();
52
53    for name in names {
54        let mut search: &str = &name;
55
56        while let Some(idx) = search.rfind('.') {
57            package_names.insert(search[0..idx].to_string());
58            search = &search[0..idx];
59        }
60    }
61
62    package_names
63}
64
65/// Resolve the filesystem path for a module.
66///
67/// Takes a path prefix, fully-qualified module name, whether the module is a package,
68/// and an optional bytecode tag to apply.
69pub fn resolve_path_for_module(
70    root: &str,
71    name: &str,
72    is_package: bool,
73    bytecode_tag: Option<&str>,
74) -> PathBuf {
75    let mut module_path = PathBuf::from(root);
76
77    let parts = name.split('.').collect::<Vec<&str>>();
78
79    // All module parts up to the final one are packages/directories.
80    for part in &parts[0..parts.len() - 1] {
81        module_path.push(*part);
82    }
83
84    // A package always exists in its own directory.
85    if is_package {
86        module_path.push(parts[parts.len() - 1]);
87    }
88
89    // If this is a bytecode module, files go in a __pycache__ directories.
90    if bytecode_tag.is_some() {
91        module_path.push("__pycache__");
92    }
93
94    // Packages get normalized to /__init__.py.
95    let basename = if is_package {
96        "__init__"
97    } else {
98        parts[parts.len() - 1]
99    };
100
101    let suffix = if let Some(tag) = bytecode_tag {
102        format!(".{}.pyc", tag)
103    } else {
104        ".py".to_string()
105    };
106
107    module_path.push(format!("{}{}", basename, suffix));
108
109    module_path
110}
111
112pub fn is_package_from_path(path: &Path) -> bool {
113    let file_name = path.file_name().unwrap().to_str().unwrap();
114    file_name.starts_with("__init__.")
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_packages_from_module_name() {
123        assert_eq!(
124            packages_from_module_name("foo.bar"),
125            ["foo".to_string()].iter().cloned().collect()
126        );
127        assert_eq!(
128            packages_from_module_name("foo.bar.baz"),
129            ["foo".to_string(), "foo.bar".to_string()]
130                .iter()
131                .cloned()
132                .collect()
133        );
134    }
135
136    #[test]
137    fn test_resolve_path_for_module() {
138        assert_eq!(
139            resolve_path_for_module(".", "foo", false, None),
140            PathBuf::from("./foo.py")
141        );
142        assert_eq!(
143            resolve_path_for_module(".", "foo", false, Some("cpython-37")),
144            PathBuf::from("./__pycache__/foo.cpython-37.pyc")
145        );
146        assert_eq!(
147            resolve_path_for_module(".", "foo", true, None),
148            PathBuf::from("./foo/__init__.py")
149        );
150        assert_eq!(
151            resolve_path_for_module(".", "foo", true, Some("cpython-37")),
152            PathBuf::from("./foo/__pycache__/__init__.cpython-37.pyc")
153        );
154        assert_eq!(
155            resolve_path_for_module(".", "foo.bar", false, None),
156            PathBuf::from("./foo/bar.py")
157        );
158        assert_eq!(
159            resolve_path_for_module(".", "foo.bar", false, Some("cpython-37")),
160            PathBuf::from("./foo/__pycache__/bar.cpython-37.pyc")
161        );
162        assert_eq!(
163            resolve_path_for_module(".", "foo.bar", true, None),
164            PathBuf::from("./foo/bar/__init__.py")
165        );
166        assert_eq!(
167            resolve_path_for_module(".", "foo.bar", true, Some("cpython-37")),
168            PathBuf::from("./foo/bar/__pycache__/__init__.cpython-37.pyc")
169        );
170        assert_eq!(
171            resolve_path_for_module(".", "foo.bar.baz", false, None),
172            PathBuf::from("./foo/bar/baz.py")
173        );
174        assert_eq!(
175            resolve_path_for_module(".", "foo.bar.baz", false, Some("cpython-37")),
176            PathBuf::from("./foo/bar/__pycache__/baz.cpython-37.pyc")
177        );
178        assert_eq!(
179            resolve_path_for_module(".", "foo.bar.baz", true, None),
180            PathBuf::from("./foo/bar/baz/__init__.py")
181        );
182        assert_eq!(
183            resolve_path_for_module(".", "foo.bar.baz", true, Some("cpython-37")),
184            PathBuf::from("./foo/bar/baz/__pycache__/__init__.cpython-37.pyc")
185        );
186    }
187}
python_packaging/module_util.rs

python_packaging/
module_util.rs
