Merge pull request #233 from Glyphack/star-imports

feat: add import stars to symbol table

parser/src/parser/parser.rs

@@ -1469,8 +1469,9 @@ impl Parser {
                 }
             }
         } else if self.at(Kind::Mul) {
-            aliases.push(self.parse_alias("*".to_string(), self.start_node()));
-            self.bump(Kind::Mul);
+            let node = self.start_node();
+            self.bump_any();
+            aliases.push(self.parse_alias("*".to_string(), node));
         } else {
             return Err(self.unexpected_token_new(
                import_node,

typechecker/src/build.rs

@@ -77,6 +77,10 @@ impl BuildManager {
                 let module = EnderpyFile::new(path, true);
                 self.modules.insert(module.path(), module);
             }
+            for (_, implicit_import) in imported_module.1.implicit_imports.iter() {
+                let module = EnderpyFile::new(&implicit_import.path, true);
+                self.modules.insert(module.path(), module);
+            }
         }
         log::debug!("Imports resolved");
         for mut module in self.modules.iter_mut() {
@@ -198,6 +202,17 @@ impl BuildManager {
                         let e = EnderpyFile::new(resolved_path, true);
                         files_to_resolve.push(e);
                     }
+
+                    for (_, implicit_import) in resolved.implicit_imports.iter() {
+                        let source = match std::fs::read_to_string(&implicit_import.path) {
+                            Ok(source) => source,
+                            Err(e) => {
+                                panic!("cannot read implicit import");
+                            }
+                        };
+                        let e = EnderpyFile::new(&implicit_import.path, true);
+                        files_to_resolve.push(e);
+                    }
                 }
             }
         }
@@ -309,4 +324,8 @@ mod tests {
         test_symbols_variables,
         "test_data/inputs/symbol_table/variables.py"
     );
+    symbol_table_test!(
+        test_symbols_import_star,
+        "test_data/inputs/import_star_test/a.py"
+    );
 }

typechecker/src/checker.rs

@@ -1,8 +1,11 @@
+use std::path::PathBuf;
+
 use ast::{Expression, Statement};
 use enderpy_python_parser as parser;
 use enderpy_python_parser::ast::{self, *};
 
 use super::{type_evaluator::TypeEvaluator, types::PythonType};
+use crate::types::ModuleRef;
 use crate::{
     ast_visitor::TraversalVisitor, diagnostic::CharacterSpan, file::EnderpyFile,
     symbol_table::SymbolTable,
@@ -164,6 +167,17 @@ impl TraversalVisitor for TypeChecker {
         for alias in _i.names.iter() {
             self.infer_name_type(&alias.name, alias.node.start, alias.node.end)
         }
+
+        // Just to show type module when modules are hovered in imports.
+        let start = _i.node.start + 5;
+        let stop = start + _i.module.len() as u32 + 1;
+        self.types.insert(Interval {
+            start,
+            stop,
+            val: PythonType::Module(ModuleRef {
+                module_path: PathBuf::new(),
+            }),
+        });
     }
 
     fn visit_if(&mut self, i: &parser::ast::If) {
@@ -669,4 +683,5 @@ mod tests {
 
     type_eval_test!(basic_types, "test_data/inputs/basic_types.py");
     type_eval_test!(basic_generics, "test_data/inputs/basic_generics.py");
+    type_eval_test!(import_star_lookup, "test_data/inputs/import_star_test/a.py");
 }

typechecker/src/semantic_analyzer.rs

@@ -346,7 +346,8 @@ impl TraversalVisitor for SemanticAnalyzer {
                 declaration_path,
                 import_from_node: None,
                 import_node: Some(i.clone()),
-                symbol_name: Some(alias.name()),
+                symbol_name: None,
+                module_name: Some(alias.name()),
                 import_result,
             });
 
@@ -361,7 +362,12 @@ impl TraversalVisitor for SemanticAnalyzer {
             None => ImportResult::not_found(),
         };
         for alias in &_i.names {
-            // TODO: make import * special to look into the other symbol table
+            if alias.name == "*" {
+                self.symbol_table
+                    .star_imports
+                    .push(module_import_result.clone());
+                continue;
+            }
             let declaration_path = DeclarationPath::new(
                 self.file.path(),
                 alias.node,
@@ -372,6 +378,7 @@ impl TraversalVisitor for SemanticAnalyzer {
                 import_from_node: Some(_i.clone()),
                 import_node: None,
                 symbol_name: Some(alias.name()),
+                module_name: None,
                 import_result: module_import_result.clone(),
             });
 

typechecker/src/symbol_table.rs

@@ -20,6 +20,7 @@ pub struct SymbolTable {
 
     pub file_path: PathBuf,
     pub scope_starts: Lapper<u32, u32>,
+    pub star_imports: Vec<ImportResult>,
 }
 
 impl SymbolTable {
@@ -38,6 +39,7 @@ impl SymbolTable {
             prev_scope_id: None,
             file_path: file_path.to_path_buf(),
             scope_starts: Lapper::new(vec![global_scope_interval]),
+            star_imports: vec![],
         }
     }
 
@@ -509,6 +511,9 @@ pub struct Alias {
     /// Name of the imported symbol in case of ImportFrom
     /// e.g. From bar import baz -> baz is the symbol name
     pub symbol_name: Option<String>,
+    /// Name of imported module in case of Import
+    /// e.g. import os.path -> os.path is the module name
+    pub module_name: Option<String>,
     /// The result of the import
     pub import_result: ImportResult,
 }
@@ -519,6 +524,7 @@ pub struct TypeAlias {
     pub type_alias_node: ast::TypeAlias,
 }
 
+#[derive(Clone, Debug)]
 pub struct LookupSymbolRequest<'a> {
     pub name: &'a str,
     pub scope: Option<u32>,
@@ -557,6 +563,9 @@ impl SymbolTableNode {
 
 impl Display for SymbolTable {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if !self.star_imports.is_empty() {
+            writeln!(f, "{:?}", self.star_imports)?;
+        }
         let mut sorted_scopes = self.scopes.iter().collect::<Vec<&SymbolTableScope>>();
         sorted_scopes.sort_by(|a, b| a.name.cmp(&b.name));
 

typechecker/src/type_evaluator.rs

@@ -16,6 +16,7 @@ use super::{
     types::{CallableType, LiteralValue, PythonType},
 };
 use crate::{
+    ruff_python_import_resolver::import_result::ImportResult,
     semantic_analyzer::get_member_access_info,
     symbol_table::{self, Class, Declaration, LookupSymbolRequest, SymbolTable, SymbolTableNode},
     types::{ClassType, TypeVar},
@@ -306,6 +307,15 @@ impl TypeEvaluator {
                             }
                         }
                     }
+                    PythonType::Module(module) => {
+                        todo!();
+                        let module_sym_table =
+                            self.get_symbol_table_of(&module.module_path).unwrap();
+                        let typ = self
+                            .infer_type_from_symbol_table(&a.attr, None, module_sym_table, Some(0))
+                            .unwrap();
+                        Ok(typ)
+                    }
                     _ => Ok(PythonType::Unknown),
                 }
             }
@@ -444,11 +454,12 @@ impl TypeEvaluator {
             name,
             symbol_table.file_path,
         );
-        let result = match symbol_table.lookup_in_scope(lookup_request) {
-            Some(symbol) => self.get_symbol_node_type(symbol),
-            None => bail!("name {name} is not defined"),
-        };
-        result
+        let result =
+            match self.lookup_in_table_and_star_imports(symbol_table, lookup_request.clone()) {
+                Some(s) => s,
+                None => bail!("name {name} is not defined"),
+            };
+        self.get_symbol_node_type(result)
     }
 
     /// Get the type of a symbol node based on declarations
@@ -537,10 +548,35 @@ impl TypeEvaluator {
                 }
             }
             Declaration::Alias(a) => {
-                let resolved = self.resolve_alias(a, symbol_table);
-                match resolved {
-                    Some(node) => self.get_symbol_node_type(node),
-                    None => Ok(PythonType::Unknown),
+                log::debug!("evaluating alias {:?}", a);
+                // when symbol is an alias that is named to that symbol return Module type
+                // e.g. from . import path as _path
+                // then type of _path is Module(path)
+                match &a.symbol_name {
+                    Some(name) => {
+                        let resolved = self.resolve_alias(name, &a.import_result, symbol_table);
+                        match resolved {
+                            Some(node) => self.get_symbol_node_type(node),
+                            None => Ok(PythonType::Unknown),
+                        }
+                    }
+                    None => {
+                        let mut found_module: Option<PythonType> = None;
+                        for resolved_path in a.import_result.resolved_paths.iter() {
+                            let Some(sym_table_alias_pointing_to) =
+                                self.get_symbol_table_of(resolved_path)
+                            else {
+                                break;
+                            };
+                            found_module = Some(PythonType::Module(crate::types::ModuleRef {
+                                module_path: sym_table_alias_pointing_to.file_path.clone(),
+                            }));
+                        }
+                        match found_module {
+                            Some(s) => Ok(s),
+                            None => Ok(PythonType::Unknown),
+                        }
+                    }
                 }
             }
             Declaration::TypeParameter(_) => Ok(PythonType::Unknown),
@@ -966,25 +1002,18 @@ impl TypeEvaluator {
         false
     }
 
-    // Follows Alias declaration and resolves it to a class declaration
+    // Follows Alias declaration and resolves it to another symbol node
     // It searches through imported symbol tables for the module alias imports
-    // and resolves the alias to the class declaration
-    // TODO: refactor all aliases and not only classes
     fn resolve_alias(
         &self,
-        a: &symbol_table::Alias,
+        // a: &symbol_table::Alias,
+        imported_symbol_name: &str,
+        import_result: &ImportResult,
         alias_symbol_table: &SymbolTable,
     ) -> Option<&symbol_table::SymbolTableNode> {
-        log::debug!("resolving alias: {:?}", a);
-        let class_name = match a.symbol_name {
-            Some(ref name) => name.clone(),
-            None => panic!("Alias {:?} has no symbol name", a.import_node),
-        };
-
-        for resolved_path in a.import_result.resolved_paths.iter() {
+        for resolved_path in import_result.resolved_paths.iter() {
             log::debug!("checking path {:?}", resolved_path);
-            // TODO: This is a hack to resolve Iterator alias in sys/__init__.pyi
-            let symbol_table_with_alias_def = if class_name == "Iterator" {
+            let symbol_table_with_alias_def = if imported_symbol_name == "Iterator" {
                 self.imported_symbol_tables
                     .iter()
                     .find(|symbol_table| symbol_table.file_path.ends_with("stdlib/typing.pyi"))
@@ -1001,12 +1030,12 @@ impl TypeEvaluator {
             // then it's cyclic and do not resolve
             if alias_symbol_table.file_path.as_path() == resolved_path {
                 log::debug!("alias resolution skipped");
-                return None;
+                continue;
             }
 
             let resolved_symbol =
                 symbol_table_with_alias_def?.lookup_in_scope(LookupSymbolRequest {
-                    name: &class_name,
+                    name: imported_symbol_name,
                     scope: None,
                 });
 
@@ -1017,6 +1046,38 @@ impl TypeEvaluator {
                 return resolved_symbol;
             }
         }
+
+        for (_, implicit_import) in import_result.implicit_imports.iter() {
+            let resolved_path = &implicit_import.path;
+            log::debug!("checking path {:?}", resolved_path);
+            // TODO: This is a hack to resolve Iterator alias in sys/__init__.pyi
+            let symbol_table_with_alias_def = if imported_symbol_name == "Iterator" {
+                self.imported_symbol_tables
+                    .iter()
+                    .find(|symbol_table| symbol_table.file_path.ends_with("stdlib/typing.pyi"))
+            } else {
+                self.get_symbol_table_of(resolved_path)
+            };
+
+            if symbol_table_with_alias_def.is_none() {
+                panic!("Symbol table not found for alias: {:?}", resolved_path);
+            }
+
+            let resolved_symbol = self.lookup_in_table_and_star_imports(
+                symbol_table_with_alias_def?,
+                LookupSymbolRequest {
+                    name: imported_symbol_name,
+                    scope: None,
+                },
+            );
+
+            log::debug!("symbols {:?}", symbol_table_with_alias_def?.global_scope());
+
+            if resolved_symbol.is_some() {
+                log::debug!("alias resolved to {:?}", resolved_symbol);
+                return resolved_symbol;
+            }
+        }
         None
     }
 
@@ -1056,4 +1117,38 @@ impl TypeEvaluator {
     ) -> PythonType {
         f_type.return_type.clone()
     }
+
+    fn lookup_in_table_and_star_imports<'a>(
+        &'a self,
+        symbol_table: &'a SymbolTable,
+        lookup_request: LookupSymbolRequest,
+    ) -> Option<&SymbolTableNode> {
+        let find_in_current_symbol_table = symbol_table.lookup_in_scope(lookup_request.clone());
+        if find_in_current_symbol_table.is_some() {
+            return find_in_current_symbol_table;
+        }
+
+        log::debug!(
+            "did not find symbol {} in symbol table, checking star imports",
+            lookup_request.name
+        );
+        // Check if there's any import * and try to find the symbol in those files
+        let mut found = None;
+        for star_import in symbol_table.star_imports.iter() {
+            log::debug!("checking star imports {:?}", star_import);
+            for resolved in star_import.resolved_paths.iter() {
+                log::debug!("checking path {:?}", resolved);
+                let star_import_sym_table = self.get_symbol_table_of(resolved);
+                let Some(sym_table) = star_import_sym_table else {
+                    panic!("symbol table of star import not found at {:?}", resolved);
+                };
+                let res = sym_table.lookup_in_scope(lookup_request.clone());
+                if res.is_some() {
+                    found = res;
+                    break;
+                }
+            }
+        }
+        found
+    }
 }