feat: support type comment

README.md

@@ -147,7 +147,7 @@ pylyzer converts Python ASTs to Erg ASTs and passes them to Erg's type checker.
 * [x] type narrowing (`is`, `isinstance`)
 * [ ] `pyi` (stub) files support
 * [ ] glob pattern file check
-* [ ] `# type: ignore` directive
+* [x] type comment (`# type: ...`)
 
 ## Join us!
 

crates/py2erg/convert.rs

@@ -21,16 +21,20 @@ use erg_compiler::erg_parser::ast::{
     VarRecordAttr, VarRecordAttrs, VarRecordPattern, VarSignature, VisModifierSpec,
 };
 use erg_compiler::erg_parser::desugar::Desugarer;
-use erg_compiler::erg_parser::token::{Token, TokenKind, COLON, DOT, EQUAL};
+use erg_compiler::erg_parser::token::{Token, TokenKind, AS, COLON, DOT, EQUAL};
 use erg_compiler::erg_parser::Parser;
 use erg_compiler::error::{CompileError, CompileErrors};
+use rustpython_ast::located::LocatedMut;
+use rustpython_ast::source_code::RandomLocator;
 use rustpython_parser::ast::located::{
     self as py_ast, Alias, Arg, Arguments, BoolOp, CmpOp, ExprConstant, Keyword, Located,
     ModModule, Operator, Stmt, String, Suite, TypeParam, UnaryOp as UnOp,
 };
+use rustpython_parser::ast::Fold;
 use rustpython_parser::source_code::{
     OneIndexed, SourceLocation as PyLocation, SourceRange as PySourceRange,
 };
+use rustpython_parser::Parse;
 
 use crate::ast_util::accessor_name;
 use crate::error::*;
@@ -278,6 +282,61 @@ impl LocalContext {
     }
 }
 
+#[derive(Debug, Default)]
+pub struct CommentStorage {
+    comments: HashMap<u32, (String, Option<py_ast::Expr>)>,
+}
+
+impl fmt::Display for CommentStorage {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for (i, (comment, expr)) in &self.comments {
+            writeln!(f, "line {i}: \"{comment}\" (expr: {})", expr.is_some())?;
+        }
+        Ok(())
+    }
+}
+
+impl CommentStorage {
+    pub fn new() -> Self {
+        Self {
+            comments: HashMap::new(),
+        }
+    }
+
+    pub fn read(&mut self, code: &str) {
+        // NOTE: This locater is meaningless.
+        let mut locater = RandomLocator::new(code);
+        for (i, line) in code.lines().enumerate() {
+            let mut split = line.split('#');
+            let _code = split.next().unwrap();
+            if let Some(comment) = split.next() {
+                let comment = comment.to_string();
+                let trimmed = comment.trim_start();
+                let expr = if trimmed.starts_with("type:") {
+                    let typ = trimmed.trim_start_matches("type:").trim();
+                    let typ = if typ == "ignore" { "Any" } else { typ };
+                    rustpython_ast::Expr::parse(typ, "<module>")
+                        .ok()
+                        .and_then(|expr| locater.fold(expr).ok())
+                } else {
+                    None
+                };
+                self.comments.insert(i as u32, (comment, expr));
+            }
+        }
+    }
+
+    /// line: 0-origin
+    pub fn get_code(&self, line: u32) -> Option<&String> {
+        self.comments.get(&line).map(|(code, _)| code)
+    }
+
+    /// line: 0-origin
+    pub fn get_type(&self, line: u32) -> Option<&py_ast::Expr> {
+        self.comments.get(&line).and_then(|(_, ty)| ty.as_ref())
+    }
+}
+
 /// AST must be converted in the following order:
 ///
 /// Params -> Block -> Signature
@@ -306,6 +365,7 @@ impl LocalContext {
 pub struct ASTConverter {
     cfg: ErgConfig,
     shadowing: ShadowingMode,
+    comments: CommentStorage,
     block_id_counter: usize,
     block_ids: Vec<usize>,
     contexts: Vec<LocalContext>,
@@ -314,10 +374,11 @@ pub struct ASTConverter {
 }
 
 impl ASTConverter {
-    pub fn new(cfg: ErgConfig, shadowing: ShadowingMode) -> Self {
+    pub fn new(cfg: ErgConfig, shadowing: ShadowingMode, comments: CommentStorage) -> Self {
         Self {
             shadowing,
             cfg,
+            comments,
             block_id_counter: 0,
             block_ids: vec![0],
             contexts: vec![LocalContext::new("<module>".into())],
@@ -2148,17 +2209,36 @@ impl ASTConverter {
                             }
                             let can_shadow = self.register_name_info(&name.id, NameKind::Variable);
                             let ident = self.convert_ident(name.id.to_string(), name.location());
+                            let t_spec = expr
+                                .ln_end()
+                                .and_then(|i| {
+                                    i.checked_sub(1)
+                                        .and_then(|line| self.comments.get_type(line))
+                                })
+                                .cloned()
+                                .map(|mut expr| {
+                                    // The range of `expr` is not correct, so we need to change it
+                                    if let py_ast::Expr::Subscript(sub) = &mut expr {
+                                        sub.range = name.range;
+                                        *sub.slice.range_mut() = name.range;
+                                        *sub.value.range_mut() = name.range;
+                                    } else {
+                                        *expr.range_mut() = name.range;
+                                    }
+                                    let t_as_expr = self.convert_expr(expr.clone());
+                                    TypeSpecWithOp::new(AS, self.convert_type_spec(expr), t_as_expr)
+                                });
                             if can_shadow.is_yes() {
                                 let block = Block::new(vec![expr]);
                                 let body = DefBody::new(EQUAL, block, DefId(0));
                                 let sig = Signature::Var(VarSignature::new(
                                     VarPattern::Ident(ident),
-                                    None,
+                                    t_spec,
                                 ));
                                 let def = Def::new(sig, body);
                                 Expr::Def(def)
                             } else {
-                                let redef = ReDef::new(Accessor::Ident(ident), None, expr);
+                                let redef = ReDef::new(Accessor::Ident(ident), t_spec, expr);
                                 Expr::ReDef(redef)
                             }
                         }

crates/pylyzer_core/analyze.rs

@@ -18,7 +18,7 @@ use erg_compiler::error::{CompileError, CompileErrors};
 use erg_compiler::module::SharedCompilerResource;
 use erg_compiler::varinfo::VarInfo;
 use erg_compiler::GenericHIRBuilder;
-use py2erg::{dump_decl_package, ShadowingMode};
+use py2erg::{dump_decl_package, CommentStorage, ShadowingMode};
 use rustpython_ast::source_code::{RandomLocator, SourceRange};
 use rustpython_ast::{Fold, ModModule};
 use rustpython_parser::{Parse, ParseErrorType};
@@ -60,13 +60,15 @@ impl ASTBuildable for SimplePythonParser {
         IncompleteParseArtifact<AST, ParserRunnerErrors>,
     > {
         let filename = self.cfg.input.filename();
+        let mut comments = CommentStorage::new();
+        comments.read(&code);
         let py_program = self.parse_py_code(code)?;
         let shadowing = if cfg!(feature = "debug") {
             ShadowingMode::Visible
         } else {
             ShadowingMode::Invisible
         };
-        let converter = py2erg::ASTConverter::new(ErgConfig::default(), shadowing);
+        let converter = py2erg::ASTConverter::new(ErgConfig::default(), shadowing, comments);
         let IncompleteArtifact {
             object: Some(erg_module),
             errors,
@@ -261,6 +263,8 @@ impl PythonAnalyzer {
     ) -> Result<CompleteArtifact, IncompleteArtifact> {
         let filename = self.cfg.input.filename();
         let parser = SimplePythonParser::new(self.cfg.copy());
+        let mut comments = CommentStorage::new();
+        comments.read(&py_code);
         let py_program = parser
             .parse_py_code(py_code)
             .map_err(|iart| IncompleteArtifact::new(None, iart.errors.into(), iart.warns.into()))?;
@@ -269,7 +273,7 @@ impl PythonAnalyzer {
         } else {
             ShadowingMode::Invisible
         };
-        let converter = py2erg::ASTConverter::new(self.cfg.copy(), shadowing);
+        let converter = py2erg::ASTConverter::new(self.cfg.copy(), shadowing, comments);
         let IncompleteArtifact {
             object: Some(erg_module),
             errors,

tests/test.rs

@@ -104,7 +104,7 @@ fn exec_warns() -> Result<(), String> {
 
 #[test]
 fn exec_typespec() -> Result<(), String> {
-    expect("tests/typespec.py", 0, 14)
+    expect("tests/typespec.py", 0, 15)
 }
 
 #[test]

tests/typespec.py

@@ -52,3 +52,9 @@ def f(x: Union[int, str, None]):
 f(1)
 f("a")
 f(None)
+
+i1 = 1 # type: int
+# ERR
+i2 = 1 # type: str
+i3 = 1 # type: ignore
+i3 + "a" # OK