rez_next_package/
python_ast_parser.rs

1//! Advanced Python AST parser for package.py files using RustPython
2
3use crate::Package;
4use rez_next_common::RezCoreError;
5use rez_next_version::Version;
6use rustpython_ast::{Constant, Expr, Stmt, Suite};
7use rustpython_parser::Parse;
8use std::collections::HashMap;
9
10/// Advanced Python AST parser for package.py files
11pub struct PythonAstParser;
12
13impl PythonAstParser {
14    /// Parse a package.py file using Python AST
15    pub fn parse_package_py(content: &str) -> Result<Package, RezCoreError> {
16        // Parse the Python code into an AST
17        let ast = Suite::parse(content, "package.py")
18            .map_err(|e| RezCoreError::PackageParse(format!("Python syntax error: {}", e)))?;
19
20        let mut package_data = PackageData::new();
21
22        // Walk through the AST and extract package information
23        for stmt in &ast {
24            Self::process_statement(stmt, &mut package_data)?;
25        }
26
27        // Convert extracted data to Package
28        Self::build_package(package_data)
29    }
30
31    /// Process a single AST statement
32    fn process_statement(stmt: &Stmt, package_data: &mut PackageData) -> Result<(), RezCoreError> {
33        match stmt {
34            Stmt::Assign(assign) => {
35                // Handle variable assignments like: name = "value"
36                if let Some(target) = assign.targets.first() {
37                    if let Expr::Name(name_expr) = target {
38                        Self::process_assignment(&name_expr.id, &assign.value, package_data)?;
39                    }
40                }
41            }
42            Stmt::FunctionDef(func_def) => {
43                // Handle function definitions like: def commands(): ...
44                if func_def.name.as_str() == "commands" {
45                    Self::process_commands_function(&func_def.body, package_data)?;
46                }
47            }
48            _ => {
49                // Ignore other statement types for now
50            }
51        }
52        Ok(())
53    }
54
55    /// Process variable assignments
56    fn process_assignment(
57        var_name: &str,
58        value: &Expr,
59        package_data: &mut PackageData,
60    ) -> Result<(), RezCoreError> {
61        match var_name {
62            "name" => {
63                package_data.name = Some(Self::extract_string_value(value)?);
64            }
65            "version" => {
66                package_data.version = Some(Self::extract_string_value(value)?);
67            }
68            "description" => {
69                package_data.description = Some(Self::extract_string_value(value)?);
70            }
71            "build_command" => {
72                package_data.build_command = Some(Self::extract_string_value(value)?);
73            }
74            "build_system" => {
75                package_data.build_system = Some(Self::extract_string_value(value)?);
76            }
77            "uuid" => {
78                package_data.uuid = Some(Self::extract_string_value(value)?);
79            }
80            "authors" => {
81                package_data.authors = Self::extract_string_list(value)?;
82            }
83            "requires" => {
84                package_data.requires = Self::extract_string_list(value)?;
85            }
86            "build_requires" => {
87                package_data.build_requires = Self::extract_string_list(value)?;
88            }
89            "private_build_requires" => {
90                package_data.private_build_requires = Self::extract_string_list(value)?;
91            }
92            "tools" => {
93                package_data.tools = Self::extract_string_list(value)?;
94            }
95            "variants" => {
96                package_data.variants = Self::extract_variants(value)?;
97            }
98            "tests" => {
99                package_data.tests = Self::extract_tests(value)?;
100            }
101            "pre_commands" => {
102                package_data.pre_commands = Some(Self::extract_string_value(value)?);
103            }
104            "post_commands" => {
105                package_data.post_commands = Some(Self::extract_string_value(value)?);
106            }
107            "pre_test_commands" => {
108                package_data.pre_test_commands = Some(Self::extract_string_value(value)?);
109            }
110            "pre_build_commands" => {
111                package_data.pre_build_commands = Some(Self::extract_string_value(value)?);
112            }
113            "requires_rez_version" => {
114                package_data.requires_rez_version = Some(Self::extract_string_value(value)?);
115            }
116            "help" => {
117                package_data.help = Some(Self::extract_string_value(value)?);
118            }
119            "relocatable" => {
120                package_data.relocatable = Self::extract_bool_value(value)?;
121            }
122            "cachable" => {
123                package_data.cachable = Self::extract_bool_value(value)?;
124            }
125            "base" => {
126                package_data.base = Some(Self::extract_string_value(value)?);
127            }
128            "hashed_variants" => {
129                package_data.hashed_variants = Self::extract_bool_value(value)?;
130            }
131            "has_plugins" => {
132                package_data.has_plugins = Self::extract_bool_value(value)?;
133            }
134            "plugin_for" => {
135                package_data.plugin_for = Self::extract_string_list(value)?;
136            }
137            "format_version" => {
138                package_data.format_version = Some(Self::extract_int_value(value)?);
139            }
140            "preprocess" => {
141                package_data.preprocess = Some(Self::extract_string_value(value)?);
142            }
143            _ => {
144                // Store unknown fields for later processing
145                package_data
146                    .extra_fields
147                    .insert(var_name.to_string(), format!("{:?}", value));
148            }
149        }
150        Ok(())
151    }
152
153    /// Extract string value from expression
154    fn extract_string_value(expr: &Expr) -> Result<String, RezCoreError> {
155        match expr {
156            Expr::Constant(constant) => match &constant.value {
157                Constant::Str(s) => Ok(s.clone()),
158                Constant::Int(i) => Ok(i.to_string()),
159                Constant::Float(f) => Ok(f.to_string()),
160                _ => Err(RezCoreError::PackageParse(format!(
161                    "Expected string/number value, got: {:?}",
162                    constant.value
163                ))),
164            },
165            _ => Err(RezCoreError::PackageParse(format!(
166                "Expected constant value, got: {:?}",
167                expr
168            ))),
169        }
170    }
171
172    /// Extract boolean value from expression
173    fn extract_bool_value(expr: &Expr) -> Result<Option<bool>, RezCoreError> {
174        match expr {
175            Expr::Constant(constant) => match &constant.value {
176                Constant::Bool(b) => Ok(Some(*b)),
177                Constant::None => Ok(None),
178                _ => Err(RezCoreError::PackageParse(format!(
179                    "Expected boolean value, got: {:?}",
180                    constant.value
181                ))),
182            },
183            _ => Err(RezCoreError::PackageParse(format!(
184                "Expected constant value, got: {:?}",
185                expr
186            ))),
187        }
188    }
189
190    /// Extract integer value from expression
191    fn extract_int_value(expr: &Expr) -> Result<i32, RezCoreError> {
192        match expr {
193            Expr::Constant(constant) => {
194                match &constant.value {
195                    Constant::Int(i) => {
196                        // Convert BigInt to i32 safely
197                        i.to_string().parse::<i32>().map_err(|e| {
198                            RezCoreError::PackageParse(format!("Integer too large for i32: {}", e))
199                        })
200                    }
201                    _ => Err(RezCoreError::PackageParse(format!(
202                        "Expected integer value, got: {:?}",
203                        constant.value
204                    ))),
205                }
206            }
207            _ => Err(RezCoreError::PackageParse(format!(
208                "Expected constant value, got: {:?}",
209                expr
210            ))),
211        }
212    }
213
214    /// Extract list of strings from expression
215    fn extract_string_list(expr: &Expr) -> Result<Vec<String>, RezCoreError> {
216        match expr {
217            Expr::List(list) => {
218                let mut result = Vec::new();
219                for elt in &list.elts {
220                    result.push(Self::extract_string_value(elt)?);
221                }
222                Ok(result)
223            }
224            Expr::Tuple(tuple) => {
225                let mut result = Vec::new();
226                for elt in &tuple.elts {
227                    result.push(Self::extract_string_value(elt)?);
228                }
229                Ok(result)
230            }
231            _ => Err(RezCoreError::PackageParse(format!(
232                "Expected list, got: {:?}",
233                expr
234            ))),
235        }
236    }
237
238    /// Extract variants (list of lists)
239    fn extract_variants(expr: &Expr) -> Result<Vec<Vec<String>>, RezCoreError> {
240        match expr {
241            Expr::List(list) => {
242                let mut result = Vec::new();
243                for elt in &list.elts {
244                    result.push(Self::extract_string_list(elt)?);
245                }
246                Ok(result)
247            }
248            _ => Err(RezCoreError::PackageParse(format!(
249                "Expected list of lists for variants, got: {:?}",
250                expr
251            ))),
252        }
253    }
254
255    /// Extract tests dictionary
256    fn extract_tests(expr: &Expr) -> Result<HashMap<String, String>, RezCoreError> {
257        match expr {
258            Expr::Dict(dict) => {
259                let mut result = HashMap::new();
260                for (key, value) in dict.keys.iter().zip(dict.values.iter()) {
261                    if let Some(key) = key {
262                        let key_str = Self::extract_string_value(key)?;
263                        let value_str = Self::extract_string_value(value)?;
264                        result.insert(key_str, value_str);
265                    }
266                }
267                Ok(result)
268            }
269            _ => Err(RezCoreError::PackageParse(format!(
270                "Expected dictionary for tests, got: {:?}",
271                expr
272            ))),
273        }
274    }
275
276    /// Process commands function
277    fn process_commands_function(
278        body: &[Stmt],
279        package_data: &mut PackageData,
280    ) -> Result<(), RezCoreError> {
281        // Extract environment variable assignments and path modifications
282        let mut commands = Vec::new();
283
284        for stmt in body {
285            if let Some(command) = Self::extract_command_from_statement(stmt)? {
286                commands.push(command);
287            }
288        }
289
290        if !commands.is_empty() {
291            package_data.commands_function = Some(commands.join("\n"));
292        }
293
294        Ok(())
295    }
296
297    /// Extract command from a statement in commands function
298    fn extract_command_from_statement(stmt: &Stmt) -> Result<Option<String>, RezCoreError> {
299        match stmt {
300            // Handle env.VAR = "value" or env.VAR.append("value")
301            Stmt::Assign(assign) => {
302                if let Some(target) = assign.targets.first() {
303                    if let Expr::Attribute(attr) = target {
304                        if let Expr::Name(name_expr) = &*attr.value {
305                            if name_expr.id.as_str() == "env" {
306                                let var_name = &attr.attr;
307                                if let Some(value) = Self::extract_string_value(&assign.value).ok()
308                                {
309                                    return Ok(Some(format!("export {}=\"{}\"", var_name, value)));
310                                }
311                            }
312                        }
313                    }
314                }
315            }
316            // Handle env.PATH.append("value") or env.VAR.prepend("value")
317            Stmt::Expr(expr_stmt) => {
318                if let Expr::Call(call) = &*expr_stmt.value {
319                    if let Expr::Attribute(attr) = &*call.func {
320                        if let Expr::Attribute(env_attr) = &*attr.value {
321                            if let Expr::Name(name_expr) = &*env_attr.value {
322                                if name_expr.id.as_str() == "env" {
323                                    let var_name = &env_attr.attr;
324                                    let method = &attr.attr;
325
326                                    if let Some(arg) = call.args.first() {
327                                        if let Ok(value) = Self::extract_string_value(arg) {
328                                            match method.as_str() {
329                                                "append" => {
330                                                    return Ok(Some(format!(
331                                                        "export {}=\"${{{}}}:{}\"",
332                                                        var_name, var_name, value
333                                                    )))
334                                                }
335                                                "prepend" => {
336                                                    return Ok(Some(format!(
337                                                        "export {}=\"{}:${{{}}}\"",
338                                                        var_name, value, var_name
339                                                    )))
340                                                }
341                                                _ => {}
342                                            }
343                                        }
344                                    }
345                                }
346                            }
347                        }
348                    }
349                }
350            }
351            _ => {}
352        }
353
354        Ok(None)
355    }
356
357    /// Build Package from extracted data
358    fn build_package(data: PackageData) -> Result<Package, RezCoreError> {
359        let name = data
360            .name
361            .ok_or_else(|| RezCoreError::PackageParse("Missing 'name' field".to_string()))?;
362
363        let mut package = Package::new(name);
364
365        // Set version
366        if let Some(version_str) = data.version {
367            package.version = Some(
368                Version::parse(&version_str)
369                    .map_err(|e| RezCoreError::PackageParse(format!("Invalid version: {}", e)))?,
370            );
371        }
372
373        // Set other fields
374        package.description = data.description;
375        package.build_command = data.build_command;
376        package.build_system = data.build_system;
377        package.pre_commands = data.pre_commands;
378        package.post_commands = data.post_commands;
379        package.pre_test_commands = data.pre_test_commands;
380        package.pre_build_commands = data.pre_build_commands;
381        package.tests = data.tests;
382        package.requires_rez_version = data.requires_rez_version;
383        package.uuid = data.uuid;
384        package.authors = data.authors;
385        package.requires = data.requires;
386        package.build_requires = data.build_requires;
387        package.private_build_requires = data.private_build_requires;
388        package.tools = data.tools;
389        package.variants = data.variants;
390        package.help = data.help;
391        package.relocatable = data.relocatable;
392        package.cachable = data.cachable;
393        package.commands = data.commands_function;
394
395        // Set new fields for complete rez compatibility
396        package.base = data.base;
397        package.hashed_variants = data.hashed_variants;
398        package.has_plugins = data.has_plugins;
399        package.plugin_for = data.plugin_for;
400        package.format_version = data.format_version;
401        package.preprocess = data.preprocess;
402
403        // Validate the package
404        package.validate()?;
405
406        Ok(package)
407    }
408}
409
410/// Intermediate data structure for collecting package information
411#[derive(Debug, Default)]
412struct PackageData {
413    name: Option<String>,
414    version: Option<String>,
415    description: Option<String>,
416    build_command: Option<String>,
417    build_system: Option<String>,
418    pre_commands: Option<String>,
419    post_commands: Option<String>,
420    pre_test_commands: Option<String>,
421    pre_build_commands: Option<String>,
422    tests: HashMap<String, String>,
423    requires_rez_version: Option<String>,
424    uuid: Option<String>,
425    authors: Vec<String>,
426    requires: Vec<String>,
427    build_requires: Vec<String>,
428    private_build_requires: Vec<String>,
429    tools: Vec<String>,
430    variants: Vec<Vec<String>>,
431    help: Option<String>,
432    relocatable: Option<bool>,
433    cachable: Option<bool>,
434    commands_function: Option<String>,
435    extra_fields: HashMap<String, String>,
436    // New fields for complete rez compatibility
437    base: Option<String>,
438    hashed_variants: Option<bool>,
439    has_plugins: Option<bool>,
440    plugin_for: Vec<String>,
441    format_version: Option<i32>,
442    preprocess: Option<String>,
443}
444
445impl PackageData {
446    fn new() -> Self {
447        Self::default()
448    }
449}
450
451#[cfg(test)]
452mod tests {
453    use super::*;
454
455    #[test]
456    fn test_parse_package_with_new_fields() {
457        let package_py_content = r#"
458name = "test_package"
459version = "1.0.0"
460description = "Test package with new fields"
461base = "base_package"
462hashed_variants = True
463has_plugins = True
464plugin_for = ["maya", "nuke"]
465format_version = 2
466preprocess = "some_preprocess_function"
467"#;
468
469        let result = PythonAstParser::parse_package_py(package_py_content);
470        assert!(
471            result.is_ok(),
472            "Failed to parse package.py: {:?}",
473            result.err()
474        );
475
476        let package = result.unwrap();
477        assert_eq!(package.name, "test_package");
478        assert_eq!(package.base, Some("base_package".to_string()));
479        assert_eq!(package.hashed_variants, Some(true));
480        assert_eq!(package.has_plugins, Some(true));
481        assert_eq!(package.plugin_for, vec!["maya", "nuke"]);
482        assert_eq!(package.format_version, Some(2));
483        assert_eq!(
484            package.preprocess,
485            Some("some_preprocess_function".to_string())
486        );
487    }
488
489    #[test]
490    fn test_parse_package_with_false_boolean_fields() {
491        let package_py_content = r#"
492name = "test_package"
493version = "1.0.0"
494hashed_variants = False
495has_plugins = False
496"#;
497
498        let result = PythonAstParser::parse_package_py(package_py_content);
499        assert!(
500            result.is_ok(),
501            "Failed to parse package.py: {:?}",
502            result.err()
503        );
504
505        let package = result.unwrap();
506        assert_eq!(package.hashed_variants, Some(false));
507        assert_eq!(package.has_plugins, Some(false));
508    }
509}