dissolve_python/
type_introspection_context.rs

1use anyhow::Result;
2use std::cell::RefCell;
3use std::path::Path;
4use std::rc::Rc;
5
6use crate::mypy_lsp::MypyTypeIntrospector;
7use crate::pyright_lsp::{PyrightLspClient, PyrightLspClientTrait};
8use crate::types::TypeIntrospectionMethod;
9
10/// Context that holds type introspection clients across multiple file migrations
11pub struct TypeIntrospectionContext {
12    method: TypeIntrospectionMethod,
13    pyright_client: Option<Rc<RefCell<Box<dyn PyrightLspClientTrait>>>>,
14    mypy_client: Option<Rc<RefCell<MypyTypeIntrospector>>>,
15    file_versions: std::collections::HashMap<String, i32>,
16    is_shutdown: bool,
17}
18
19impl TypeIntrospectionContext {
20    /// Create a new type introspection context
21    pub fn new(method: TypeIntrospectionMethod) -> Result<Self> {
22        Self::new_with_workspace(method, None)
23    }
24
25    /// Create a new type introspection context with a specific workspace root
26    pub fn new_with_workspace(
27        method: TypeIntrospectionMethod,
28        workspace_root: Option<&str>,
29    ) -> Result<Self> {
30        // Allow overriding type introspection method via environment variable for testing
31        #[cfg(test)]
32        let method = if let Ok(forced_method) = std::env::var("DISSOLVE_FORCE_TYPE_INTROSPECTION") {
33            match forced_method.to_lowercase().as_str() {
34                "dmypy" | "mypy" => TypeIntrospectionMethod::MypyDaemon,
35                "pyright" | "pyrightlsp" => TypeIntrospectionMethod::PyrightLsp,
36                "pyright,dmypy" | "pyright,mypy" | "pyrightfallback" => {
37                    TypeIntrospectionMethod::PyrightWithMypyFallback
38                }
39                _ => {
40                    eprintln!("Warning: Unknown DISSOLVE_FORCE_TYPE_INTROSPECTION value '{}', using original method", forced_method);
41                    method
42                }
43            }
44        } else {
45            method
46        };
47
48        let (pyright_client, mypy_client) = match method {
49            TypeIntrospectionMethod::PyrightLsp => {
50                // Always create a fresh client to avoid isolation issues
51                let client: Box<dyn PyrightLspClientTrait> =
52                    Box::new(PyrightLspClient::new(workspace_root)?);
53                (Some(Rc::new(RefCell::new(client))), None)
54            }
55            TypeIntrospectionMethod::MypyDaemon => {
56                let client = MypyTypeIntrospector::new(workspace_root)
57                    .map_err(|e| anyhow::anyhow!("Failed to create mypy client: {}", e))?;
58                (None, Some(Rc::new(RefCell::new(client))))
59            }
60            TypeIntrospectionMethod::PyrightWithMypyFallback => {
61                // Create fresh clients to avoid isolation issues
62                let pyright = match PyrightLspClient::new(workspace_root) {
63                    Ok(client) => {
64                        let client: Box<dyn PyrightLspClientTrait> = Box::new(client);
65                        Some(Rc::new(RefCell::new(client)))
66                    }
67                    Err(_) => None,
68                };
69                let mypy = match MypyTypeIntrospector::new(workspace_root) {
70                    Ok(client) => Some(Rc::new(RefCell::new(client))),
71                    Err(_) => None,
72                };
73                if pyright.is_none() && mypy.is_none() {
74                    return Err(anyhow::anyhow!(
75                        "Failed to initialize any type introspection client"
76                    ));
77                }
78                (pyright, mypy)
79            }
80        };
81
82        Ok(Self {
83            method,
84            pyright_client,
85            mypy_client,
86            file_versions: std::collections::HashMap::new(),
87            is_shutdown: false,
88        })
89    }
90
91    /// Get the type introspection method
92    pub fn method(&self) -> TypeIntrospectionMethod {
93        self.method
94    }
95
96    /// Get a clone of the pyright client if available
97    pub fn pyright_client(&self) -> Option<Rc<RefCell<Box<dyn PyrightLspClientTrait>>>> {
98        self.pyright_client.as_ref().map(|rc| rc.clone())
99    }
100
101    /// Get a clone of the mypy client if available
102    pub fn mypy_client(&self) -> Option<Rc<RefCell<MypyTypeIntrospector>>> {
103        self.mypy_client.as_ref().map(|rc| rc.clone())
104    }
105
106    /// Open a file for type introspection
107    pub fn open_file(&mut self, file_path: &Path, content: &str) -> Result<()> {
108        let path_str = file_path.to_string_lossy();
109        self.file_versions.insert(path_str.to_string(), 1);
110
111        if let Some(ref client) = self.pyright_client {
112            client.borrow_mut().open_file(&path_str, content)?;
113        }
114
115        Ok(())
116    }
117
118    /// Update a file after modifications
119    pub fn update_file(&mut self, file_path: &Path, content: &str) -> Result<()> {
120        let path_str = file_path.to_string_lossy();
121        let version = self
122            .file_versions
123            .get(path_str.as_ref())
124            .copied()
125            .unwrap_or(1)
126            + 1;
127        self.file_versions.insert(path_str.to_string(), version);
128
129        if let Some(ref client) = self.pyright_client {
130            client
131                .borrow_mut()
132                .update_file(&path_str, content, version)?;
133        }
134
135        if let Some(ref client) = self.mypy_client {
136            client
137                .borrow_mut()
138                .invalidate_file(&path_str)
139                .map_err(|e| anyhow::anyhow!("Failed to invalidate mypy cache: {}", e))?;
140        }
141
142        Ok(())
143    }
144
145    /// Check if the context has been shutdown
146    pub fn is_shutdown(&self) -> bool {
147        self.is_shutdown
148    }
149
150    /// Shutdown the clients cleanly
151    pub fn shutdown(&mut self) -> Result<()> {
152        if self.is_shutdown {
153            return Ok(());
154        }
155
156        if let Some(ref client) = self.pyright_client {
157            client.borrow_mut().shutdown()?;
158        }
159
160        if let Some(ref client) = self.mypy_client {
161            client
162                .borrow_mut()
163                .stop_daemon()
164                .map_err(|e| anyhow::anyhow!("Failed to stop mypy daemon: {}", e))?;
165        }
166
167        self.is_shutdown = true;
168        Ok(())
169    }
170}
171
172impl Drop for TypeIntrospectionContext {
173    fn drop(&mut self) {
174        // Try to shutdown cleanly, but don't panic on failure
175        let _ = self.shutdown();
176    }
177}
dissolve_python/type_introspection_context.rs

dissolve_python/
type_introspection_context.rs
