package_parser/pkgs/python/
pysetup.rs

1use tree_sitter::{Parser, Query, QueryCursor};
2
3use crate::error::SourcePkgError;
4use crate::pkgs::common::model::{Package, PackageManifest};
5
6use crate::pkgs::pyrequirements::PyRequirements;
7
8use std::path::Path;
9
10static SETUP_QUERY_STR: &str = "
11(
12	call
13    function: (identifier) @function
14    arguments: (
15    	argument_list
16        (
17        	keyword_argument
18            name: (identifier) @kw-install-requires-name
19            value: (_) @kw-value
20        )
21    )
22    (#eq? @function \"setup\")
23    (#eq? @kw-install-requires-name \"install_requires\")
24)
25";
26
27lazy_static::lazy_static! {
28    static ref SETUP_QUERY: Query = {
29        Query::new(&tree_sitter_python::language(), SETUP_QUERY_STR).unwrap()
30    };
31    static ref STRING_QUERY: Query = {
32        Query::new(&tree_sitter_python::language(), "(string (string_content) @str)").unwrap()
33    };
34}
35
36fn parse_strings(node: tree_sitter::Node, content: &[u8]) -> Vec<String> {
37    let mut query_cursor = QueryCursor::new();
38    let mut strings = vec![];
39    for m in query_cursor.matches(&STRING_QUERY, node, content) {
40        let value_node = m.captures.first().unwrap().node;
41        let value = value_node.utf8_text(content).unwrap();
42        strings.push(value.to_string());
43    }
44
45    strings
46}
47
48pub struct PySetup {}
49
50impl PySetup {
51    pub fn new() -> Self {
52        Self {}
53    }
54
55    fn parse(path: impl AsRef<Path>) -> Result<Package, SourcePkgError> {
56        let content = std::fs::read_to_string(path.as_ref())?;
57        let content_raw = content.as_bytes();
58
59        let mut parser = Parser::new();
60        parser
61            .set_language(&tree_sitter_python::language())
62            .unwrap();
63        let tree = parser.parse(content_raw, None).unwrap();
64        let root_node = tree.root_node();
65
66        let mut query_cursor = QueryCursor::new();
67        for m in query_cursor.matches(&SETUP_QUERY, root_node, content_raw) {
68            let value_node = m.captures.get(2).unwrap().node;
69
70            let packages = match value_node.grammar_name() {
71                "list" => parse_strings(value_node, content_raw),
72                "identifier" => {
73                    let var_name = value_node.utf8_text(content_raw).unwrap();
74                    let query = Query::new(
75                        &tree_sitter_python::language(),
76                        &format!(
77                            "
78                        (
79                            assignment
80                            left: (identifier) @name
81                            right: (list) @value
82                            (#eq? @name \"{}\")
83                        )
84                        ",
85                            var_name
86                        ),
87                    )
88                    .unwrap();
89
90                    let mut query_cursor = QueryCursor::new();
91                    if let Some(m) = query_cursor.matches(&query, root_node, content_raw).next() {
92                        let value_node = m.captures.get(1).unwrap().node;
93                        parse_strings(value_node, content_raw)
94                    } else {
95                        continue;
96                    }
97                }
98                _ => continue,
99            };
100
101            if packages.is_empty() {
102                continue;
103            }
104
105            let mut requirements = packages.join("\n");
106            requirements.push('\n');
107            let parsed = PyRequirements::parse_requirement_content(&requirements)?;
108
109            dbg!(parsed);
110        }
111
112        let package = Package::default();
113        Ok(package)
114    }
115}
116
117#[async_trait::async_trait]
118impl PackageManifest for PySetup {
119    fn get_name(&self) -> String {
120        "pypi".to_string()
121    }
122
123    async fn recognize(&self, path: &Path) -> Result<Package, SourcePkgError> {
124        Self::parse(path)
125    }
126
127    fn file_name_patterns(&self) -> &'static [&'static str] {
128        &["setup.py", "*setup.py", "setup*.py"]
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn test_parse_pysetup_1() {
138        let filepath = Path::new(concat!(
139            env!("CARGO_MANIFEST_DIR"),
140            "/testdata/pypi/setup.py/simple-setup.py"
141        ));
142
143        let p = PySetup::parse(filepath).unwrap();
144        println!("{:?}", p);
145    }
146
147    #[test]
148    fn test_parse_pysetup_2() {
149        let filepath = Path::new(concat!(
150            env!("CARGO_MANIFEST_DIR"),
151            "/testdata/pypi/setup.py/pipdeptree_setup.py"
152        ));
153
154        let p = PySetup::parse(filepath).unwrap();
155        println!("{:?}", p);
156    }
157}
package_parser/pkgs/python/pysetup.rs

package_parser/pkgs/python/
pysetup.rs
