From 855d4630119e8f18ee2b19eeafb9ffc1094de0cb Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Tue, 13 Dec 2022 22:38:20 +0900
Subject: [PATCH] Add `trim_end_proc_mark`

---
 compiler/erg_compiler/lower.rs | 10 ++++++++--
 compiler/erg_parser/ast.rs     | 11 +++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/compiler/erg_compiler/lower.rs b/compiler/erg_compiler/lower.rs
index 6bf7cf80..84395536 100644
--- a/compiler/erg_compiler/lower.rs
+++ b/compiler/erg_compiler/lower.rs
@@ -1909,7 +1909,10 @@ impl ASTLowerer {
         let is_instance_ascription = tasc.is_instance_ascription();
         let mut dummy_tv_cache = TyVarCache::new(self.ctx.level, &self.ctx);
         match *tasc.expr {
-            ast::Expr::Accessor(ast::Accessor::Ident(ident)) => {
+            ast::Expr::Accessor(ast::Accessor::Ident(mut ident)) => {
+                if self.cfg.pylyzer_mode {
+                    ident.trim_end_proc_mark();
+                }
                 let py_name = Str::rc(ident.inspect().trim_end_matches('!'));
                 let t = self.ctx.instantiate_typespec(
                     &tasc.t_spec,
@@ -1932,7 +1935,10 @@ impl ASTLowerer {
                 let ident = hir::Identifier::new(ident.dot, ident.name, None, vi);
                 Ok(hir::Expr::Accessor(hir::Accessor::Ident(ident)).type_asc(tasc.t_spec))
             }
-            ast::Expr::Accessor(ast::Accessor::Attr(attr)) => {
+            ast::Expr::Accessor(ast::Accessor::Attr(mut attr)) => {
+                if self.cfg.pylyzer_mode {
+                    attr.ident.trim_end_proc_mark();
+                }
                 let py_name = Str::rc(attr.ident.inspect().trim_end_matches('!'));
                 let t = self.ctx.instantiate_typespec(
                     &tasc.t_spec,
diff --git a/compiler/erg_parser/ast.rs b/compiler/erg_parser/ast.rs
index 3b322b7e..8b346e10 100644
--- a/compiler/erg_parser/ast.rs
+++ b/compiler/erg_parser/ast.rs
@@ -2295,6 +2295,13 @@ impl VarName {
     pub const fn inspect(&self) -> &Str {
         &self.0.content
     }
+
+    /// Remove `!` from the end of the identifier.
+    /// Procedures defined in `d.er` automatically register the name without `!` as `py_name`.
+    /// This method is for undoing it (e.g. pylyzer-mode)
+    pub fn trim_end_proc_mark(&mut self) {
+        let _ = self.0.content.trim_end_matches('!');
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -2378,6 +2385,10 @@ impl Identifier {
     pub fn is_procedural(&self) -> bool {
         self.name.is_procedural()
     }
+
+    pub fn trim_end_proc_mark(&mut self) {
+        self.name.trim_end_proc_mark();
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]