Skip to main content

dissolve_python/
migrate_stub.rs

1// Implementation of migrate_ruff functions using rustpython-parser
2use crate::checker::CheckResult;
3use crate::core::ReplaceInfo;
4use crate::ruff_parser::apply_replacements;
5use crate::rustpython_visitor::RustPythonFunctionCallReplacer;
6use crate::stub_collector::RuffDeprecatedFunctionCollector;
7use anyhow::Result;
8use std::collections::HashMap;
9
10pub fn migrate_file(
11    source: &str,
12    module_name: &str,
13    file_path: &std::path::Path,
14    type_context: &mut crate::type_introspection_context::TypeIntrospectionContext,
15    mut all_replacements: HashMap<String, ReplaceInfo>,
16    dependency_inheritance_map: HashMap<String, Vec<String>>,
17) -> Result<String> {
18    // Open the file in the type introspection context so we can query types
19    type_context.open_file(file_path, source)?;
20
21    // Always collect from source to get inheritance information
22    let collector = RuffDeprecatedFunctionCollector::new(module_name.to_string(), Some(file_path));
23    let builtins = collector.builtins().clone();
24    let collector_result = collector.collect_from_source(source.to_string())?;
25
26    // Merge provided replacements with ones collected from the source file
27    // Source file replacements take priority over dependency replacements
28    for (key, value) in collector_result.replacements {
29        all_replacements.insert(key, value);
30    }
31
32    // Parse source with rustpython
33    let ast = rustpython_parser::parse(source, rustpython_parser::Mode::Module, "<test>")
34        .map_err(|e| anyhow::anyhow!("Failed to parse: {:?}", e))?;
35
36    // Merge inheritance maps
37    let mut merged_inheritance_map = collector_result.inheritance_map;
38    merged_inheritance_map.extend(dependency_inheritance_map);
39
40    // Find and replace calls using rustpython visitor
41    let mut replacer = RustPythonFunctionCallReplacer::new_with_context(
42        all_replacements,
43        type_context,
44        file_path.to_string_lossy().into_owned(),
45        module_name.to_string(),
46        source.to_string(),
47        merged_inheritance_map,
48        builtins,
49    )?;
50
51    // Visit the AST to find replacements
52    if let rustpython_ast::Mod::Module(module) = ast {
53        replacer.visit_module(&module.body);
54    }
55
56    let replacements = replacer.get_replacements();
57
58    // Apply replacements
59    tracing::debug!("Applying {} replacements", replacements.len());
60    for (range, replacement) in &replacements {
61        if range.start < source.len() && range.end <= source.len() {
62            let original = &source[range.clone()];
63            tracing::debug!("Replacing '{}' with '{}'", original, replacement);
64        } else {
65            tracing::warn!(
66                "Invalid range {:?} for source length {}",
67                range,
68                source.len()
69            );
70        }
71    }
72    let migrated_source = apply_replacements(source, replacements.clone());
73
74    // Try to parse the migrated source to verify it's valid
75    if let Err(e) =
76        rustpython_parser::parse(&migrated_source, rustpython_parser::Mode::Module, "<test>")
77    {
78        tracing::error!("Generated invalid Python: {}", e);
79        tracing::error!("Migrated source:\n{}", migrated_source);
80    }
81
82    // Update the file in type introspection context if changes were made
83    if !replacements.is_empty() {
84        type_context.update_file(file_path, &migrated_source)?;
85    }
86
87    Ok(migrated_source)
88}
89
90pub fn migrate_file_interactive(
91    source: &str,
92    module_name: &str,
93    file_path: &std::path::Path,
94    type_context: &mut crate::type_introspection_context::TypeIntrospectionContext,
95    all_replacements: HashMap<String, ReplaceInfo>,
96    inheritance_map: HashMap<String, Vec<String>>,
97) -> Result<String> {
98    // For now, just use non-interactive version
99    migrate_file(
100        source,
101        module_name,
102        file_path,
103        type_context,
104        all_replacements,
105        inheritance_map,
106    )
107}
108
109pub fn check_file(
110    source: &str,
111    module_name: &str,
112    file_path: &std::path::Path,
113) -> Result<CheckResult> {
114    let collector = RuffDeprecatedFunctionCollector::new(module_name.to_string(), Some(file_path));
115    let result = collector.collect_from_source(source.to_string())?;
116
117    let mut check_result = CheckResult::new();
118
119    // Add all found functions to checked_functions
120    for func_name in result.replacements.keys() {
121        check_result.add_checked_function(func_name.clone());
122    }
123
124    // Also add unreplaceable functions to checked_functions
125    for func_name in result.unreplaceable.keys() {
126        check_result.add_checked_function(func_name.clone());
127    }
128
129    // Add errors for unreplaceable functions
130    for (func_name, unreplaceable) in result.unreplaceable {
131        check_result.add_error(format!(
132            "Function '{}' cannot be replaced: {:?}",
133            func_name, unreplaceable.reason
134        ));
135    }
136
137    Ok(check_result)
138}