diff --git a/README.md b/README.md index c9442db..01cceff 100644 --- a/README.md +++ b/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! diff --git a/crates/py2erg/convert.rs b/crates/py2erg/convert.rs index d068ce7..e0c41be 100644 --- a/crates/py2erg/convert.rs +++ b/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)>, +} + +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, "") + .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, contexts: Vec, @@ -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("".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) } } diff --git a/crates/pylyzer_core/analyze.rs b/crates/pylyzer_core/analyze.rs index 9948c6a..34fb2de 100644 --- a/crates/pylyzer_core/analyze.rs +++ b/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, > { 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 { 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, diff --git a/tests/test.rs b/tests/test.rs index 227cb4a..df1d6fd 100644 --- a/tests/test.rs +++ b/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] diff --git a/tests/typespec.py b/tests/typespec.py index 3404d0e..a9a86e6 100644 --- a/tests/typespec.py +++ b/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