feat: support `TypeVar` & type parameter syntax

User class generics are not yet supported

Cargo.lock

@@ -139,9 +139,9 @@ checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
 [[package]]
 name = "els"
-version = "0.1.54-nightly.2"
+version = "0.1.54-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "19d774d511ed129b8438a9633b17f3303768fa7a3941372c1650e61a7bc7b00c"
+checksum = "cfe2679dec933507f4d9abfa4e6468d3b312790e933b3269ce2a2a2bc910bd3d"
 dependencies = [
  "erg_common",
  "erg_compiler",
@@ -153,9 +153,9 @@ dependencies = [
 
 [[package]]
 name = "erg_common"
-version = "0.6.42-nightly.2"
+version = "0.6.42-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b084d80afdb59d10d300595b2860868e52d1a0a72ad98ac995c9f5abfba9acd8"
+checksum = "1c5582717e4cd56c2015263641441c7f3c19ae030bd804565f0b93ad0a8c536e"
 dependencies = [
  "backtrace-on-stack-overflow",
  "erg_proc_macros",
@@ -165,9 +165,9 @@ dependencies = [
 
 [[package]]
 name = "erg_compiler"
-version = "0.6.42-nightly.2"
+version = "0.6.42-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "805b009c668055c8d72d11f8b26853c9a65a2744548909beda8b32bd4058a375"
+checksum = "299406202ab6dfe28be95c3d6ceae19d9821624ef666a978b42112f658b528c4"
 dependencies = [
  "erg_common",
  "erg_parser",
@@ -175,9 +175,9 @@ dependencies = [
 
 [[package]]
 name = "erg_parser"
-version = "0.6.42-nightly.2"
+version = "0.6.42-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ec4f2cc69904baae639ff154323d1cc0e100ca2785bfb6d8c392c88e0560771"
+checksum = "cc13d9b7c3342e1a4ccd4159e39c99769a30733ac6cf44c22593d0aecb169af9"
 dependencies = [
  "erg_common",
  "erg_proc_macros",
@@ -186,9 +186,9 @@ dependencies = [
 
 [[package]]
 name = "erg_proc_macros"
-version = "0.6.42-nightly.2"
+version = "0.6.42-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6234abaef2fed929391add7520409890b2b7ee7029f2a5dcb9c2f4905bb7556b"
+checksum = "b8506a5228f462df55923b4a6ca8617ce90a84b52a4b603cedf3dda6beb3243f"
 dependencies = [
  "quote",
  "syn 1.0.109",

Cargo.toml

@@ -22,9 +22,9 @@ edition = "2021"
 repository = "https://github.com/mtshiba/pylyzer"
 
 [workspace.dependencies]
-erg_common = { version = "0.6.42-nightly.2", features = ["py_compat", "els"] }
+erg_common = { version = "0.6.42-nightly.3", features = ["py_compat", "els"] }
-erg_compiler = { version = "0.6.42-nightly.2", features = ["py_compat", "els"] }
+erg_compiler = { version = "0.6.42-nightly.3", features = ["py_compat", "els"] }
-els = { version = "0.1.54-nightly.2", features = ["py_compat"] }
+els = { version = "0.1.54-nightly.3", features = ["py_compat"] }
 # rustpython-parser = { version = "0.3.0", features = ["all-nodes-with-ranges", "location"] }
 # rustpython-ast = { version = "0.3.0", features = ["all-nodes-with-ranges", "location"] }
 rustpython-parser = { git = "https://github.com/RustPython/Parser", version = "0.4.0", features = ["all-nodes-with-ranges", "location"] }

crates/py2erg/convert.rs

@@ -1,3 +1,4 @@
+use core::fmt;
 use std::path::Path;
 
 use erg_common::config::ErgConfig;
@@ -14,9 +15,9 @@ use erg_compiler::erg_parser::ast::{
     LambdaSignature, List, ListComprehension, Literal, Methods, Module, NonDefaultParamSignature,
     NormalDict, NormalList, NormalRecord, NormalSet, NormalTuple, ParamPattern, ParamTySpec,
     Params, PosArg, PreDeclTypeSpec, ReDef, Record, RecordAttrs, Set, SetComprehension, Signature,
-    SubrSignature, SubrTypeSpec, Tuple, TupleTypeSpec, TypeAscription, TypeBoundSpecs, TypeSpec,
+    SubrSignature, SubrTypeSpec, Tuple, TupleTypeSpec, TypeAscription, TypeBoundSpec,
-    TypeSpecWithOp, UnaryOp, VarName, VarPattern, VarRecordAttr, VarRecordAttrs, VarRecordPattern,
+    TypeBoundSpecs, TypeSpec, TypeSpecWithOp, UnaryOp, VarName, VarPattern, VarRecordAttr,
-    VarSignature, VisModifierSpec,
+    VarRecordAttrs, VarRecordPattern, VarSignature, VisModifierSpec,
 };
 use erg_compiler::erg_parser::desugar::Desugarer;
 use erg_compiler::erg_parser::token::{Token, TokenKind, COLON, DOT, EQUAL};
@@ -24,7 +25,7 @@ use erg_compiler::erg_parser::Parser;
 use erg_compiler::error::{CompileError, CompileErrors};
 use rustpython_parser::ast::located::{
     self as py_ast, Alias, Arg, Arguments, BoolOp, CmpOp, ExprConstant, Keyword, Located,
-    ModModule, Operator, Stmt, String, Suite, UnaryOp as UnOp,
+    ModModule, Operator, Stmt, String, Suite, TypeParam, UnaryOp as UnOp,
 };
 use rustpython_parser::source_code::{
     OneIndexed, SourceLocation as PyLocation, SourceRange as PySourceRange,
@@ -217,6 +218,49 @@ pub enum ShadowingMode {
     Visible,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct TypeVarInfo {
+    name: String,
+    bound: Option<Expr>,
+}
+
+impl fmt::Display for TypeVarInfo {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if let Some(bound) = &self.bound {
+            write!(f, "TypeVarInfo({} bound={})", self.name, bound)
+        } else {
+            write!(f, "TypeVarInfo({})", self.name)
+        }
+    }
+}
+
+impl TypeVarInfo {
+    pub const fn new(name: String, bound: Option<Expr>) -> Self {
+        Self { name, bound }
+    }
+}
+
+#[derive(Debug)]
+pub struct LocalContext {
+    name: String,
+    /// Erg does not allow variables to be defined multiple times, so rename them using this
+    names: HashMap<String, NameInfo>,
+    type_vars: HashMap<String, TypeVarInfo>,
+    // e.g. def id(x: T) -> T: ... => appeared_types = {T}
+    appeared_type_names: HashSet<String>,
+}
+
+impl LocalContext {
+    pub fn new(name: String) -> Self {
+        Self {
+            name,
+            names: HashMap::new(),
+            type_vars: HashMap::new(),
+            appeared_type_names: HashSet::new(),
+        }
+    }
+}
+
 /// AST must be converted in the following order:
 ///
 /// Params -> Block -> Signature
@@ -245,11 +289,9 @@ pub enum ShadowingMode {
 pub struct ASTConverter {
     cfg: ErgConfig,
     shadowing: ShadowingMode,
-    namespace: Vec<String>,
     block_id_counter: usize,
     block_ids: Vec<usize>,
-    /// Erg does not allow variables to be defined multiple times, so rename them using this
+    contexts: Vec<LocalContext>,
-    names: Vec<HashMap<String, NameInfo>>,
     warns: CompileErrors,
     errs: CompileErrors,
 }
@@ -259,18 +301,17 @@ impl ASTConverter {
         Self {
             shadowing,
             cfg,
-            namespace: vec![String::from("<module>")],
             block_id_counter: 0,
             block_ids: vec![0],
-            names: vec![HashMap::new()],
+            contexts: vec![LocalContext::new("<module>".into())],
             warns: CompileErrors::empty(),
             errs: CompileErrors::empty(),
         }
     }
 
     fn get_name(&self, name: &str) -> Option<&NameInfo> {
-        for ns in self.names.iter().rev() {
+        for ctx in self.contexts.iter().rev() {
-            if let Some(ni) = ns.get(name) {
+            if let Some(ni) = ctx.names.get(name) {
                 return Some(ni);
             }
         }
@@ -278,38 +319,67 @@ impl ASTConverter {
     }
 
     fn get_mut_name(&mut self, name: &str) -> Option<&mut NameInfo> {
-        for ns in self.names.iter_mut().rev() {
+        for ctx in self.contexts.iter_mut().rev() {
-            if let Some(ni) = ns.get_mut(name) {
+            if let Some(ni) = ctx.names.get_mut(name) {
                 return Some(ni);
             }
         }
         None
     }
 
+    fn get_type_var(&self, name: &str) -> Option<&TypeVarInfo> {
+        for ctx in self.contexts.iter().rev() {
+            if let Some(tv) = ctx.type_vars.get(name) {
+                return Some(tv);
+            }
+        }
+        None
+    }
+
     fn define_name(&mut self, name: String, info: NameInfo) {
-        self.names.last_mut().unwrap().insert(name, info);
+        self.contexts.last_mut().unwrap().names.insert(name, info);
     }
 
     fn declare_name(&mut self, name: String, info: NameInfo) {
-        self.names.first_mut().unwrap().insert(name, info);
+        self.contexts.first_mut().unwrap().names.insert(name, info);
+    }
+
+    fn define_type_var(&mut self, name: String, info: TypeVarInfo) {
+        self.contexts
+            .last_mut()
+            .unwrap()
+            .type_vars
+            .insert(name, info);
     }
 
     fn grow(&mut self, namespace: String) {
-        self.namespace.push(namespace);
+        self.contexts.push(LocalContext::new(namespace));
-        self.names.push(HashMap::new());
     }
 
     fn pop(&mut self) {
-        self.namespace.pop();
+        self.contexts.pop();
-        self.names.pop();
     }
 
     fn cur_block_id(&self) -> usize {
         *self.block_ids.last().unwrap()
     }
 
+    /// foo.bar.baz
     fn cur_namespace(&self) -> String {
-        self.namespace.join(".")
+        self.contexts
+            .iter()
+            .map(|ctx| &ctx.name[..])
+            .collect::<Vec<_>>()
+            .join(".")
+    }
+
+    // baz
+    fn cur_name(&self) -> &str {
+        &self.contexts.last().unwrap().name
+    }
+
+    fn cur_appeared_type_names(&self) -> &HashSet<String> {
+        &self.contexts.last().unwrap().appeared_type_names
     }
 
     fn register_name_info(&mut self, name: &str, kind: NameKind) -> CanShadow {
@@ -869,6 +939,11 @@ impl ASTConverter {
         #[allow(clippy::collapsible_match)]
         match expr {
             py_ast::Expr::Name(name) => {
+                self.contexts
+                    .last_mut()
+                    .unwrap()
+                    .appeared_type_names
+                    .insert(name.id.to_string());
                 self.convert_ident_type_spec(name.id.to_string(), name.location())
             }
             py_ast::Expr::Constant(cons) => {
@@ -1306,7 +1381,7 @@ impl ASTConverter {
                 self.errs.push(self_not_found_error(
                     self.cfg.input.clone(),
                     subr.loc(),
-                    self.namespace.join("."),
+                    self.cur_namespace(),
                 ));
                 Some(())
             }
@@ -1314,7 +1389,7 @@ impl ASTConverter {
                 self.errs.push(init_var_error(
                     self.cfg.input.clone(),
                     var.loc(),
-                    self.namespace.join("."),
+                    self.cur_namespace(),
                 ));
                 None
             }
@@ -1418,13 +1493,17 @@ impl ASTConverter {
         );
         let class_ident = Identifier::public_with_line(
             DOT,
-            self.namespace.last().unwrap().into(),
+            self.cur_name().to_string().into(),
             sig.ln_begin().unwrap_or(0),
         );
         let class_ident_expr = Expr::Accessor(Accessor::Ident(class_ident.clone()));
         let class_spec = TypeSpecWithOp::new(COLON, TypeSpec::mono(class_ident), class_ident_expr);
         let mut params = sig.params.clone();
-        if params.non_defaults.first().is_some_and(|param| param.inspect().map(|s| &s[..]) == Some("self")) {
+        if params
+            .non_defaults
+            .first()
+            .is_some_and(|param| param.inspect().map(|s| &s[..]) == Some("self"))
+        {
             params.non_defaults.remove(0);
         }
         let sig = Signature::Subr(SubrSignature::new(
@@ -1449,7 +1528,7 @@ impl ASTConverter {
         );
         let params = Params::empty();
         let class_ident =
-            Identifier::public_with_line(DOT, self.namespace.last().unwrap().into(), line as u32);
+            Identifier::public_with_line(DOT, self.cur_name().to_string().into(), line as u32);
         let class_ident_expr = Expr::Accessor(Accessor::Ident(class_ident.clone()));
         let class_spec = TypeSpecWithOp::new(COLON, TypeSpec::mono(class_ident), class_ident_expr);
         let sig = Signature::Subr(SubrSignature::new(
@@ -1566,15 +1645,49 @@ impl ASTConverter {
         (base_type, vec![methods])
     }
 
-    fn convert_funcdef(
+    fn get_type_bounds(&mut self, type_params: Vec<TypeParam>) -> TypeBoundSpecs {
-        &mut self,
+        let mut bounds = TypeBoundSpecs::empty();
-        name: String,
+        if type_params.is_empty() {
-        params: Arguments,
+            for ty in self.cur_appeared_type_names() {
-        body: Vec<py_ast::Stmt>,
+                let name = VarName::from_str(ty.clone().into());
-        decorator_list: Vec<py_ast::Expr>,
+                let op = Token::dummy(TokenKind::SubtypeOf, "<:");
-        returns: Option<py_ast::Expr>,
+                if let Some(tv_info) = self.get_type_var(ty) {
-        range: PySourceRange,
+                    let bound = if let Some(bound) = &tv_info.bound {
-    ) -> Expr {
+                        let t_spec = Parser::expr_to_type_spec(bound.clone())
+                            .unwrap_or(TypeSpec::Infer(name.token().clone()));
+                        let spec = TypeSpecWithOp::new(op, t_spec, bound.clone());
+                        TypeBoundSpec::non_default(name, spec)
+                    } else {
+                        TypeBoundSpec::Omitted(name)
+                    };
+                    bounds.push(bound);
+                }
+            }
+        }
+        for tp in type_params {
+            // TODO:
+            let Some(tv) = tp.as_type_var() else {
+                continue;
+            };
+            let name = VarName::from_str(tv.name.to_string().into());
+            let spec = if let Some(bound) = &tv.bound {
+                let op = Token::dummy(TokenKind::SubtypeOf, "<:");
+                let spec = self.convert_type_spec(*bound.clone());
+                let expr = self.convert_expr(*bound.clone());
+                let spec = TypeSpecWithOp::new(op, spec, expr);
+                TypeBoundSpec::non_default(name, spec)
+            } else {
+                TypeBoundSpec::Omitted(name)
+            };
+            bounds.push(spec);
+        }
+        bounds
+    }
+
+    fn convert_funcdef(&mut self, func_def: py_ast::StmtFunctionDef) -> Expr {
+        let name = func_def.name.to_string();
+        let params = *func_def.args;
+        let returns = func_def.returns.map(|x| *x);
         // if reassigning of a function referenced by other functions is occurred, it is an error
         if self.get_name(&name).is_some_and(|info| {
             info.defined_times > 0
@@ -1583,15 +1696,16 @@ impl ASTConverter {
         }) {
             let err = reassign_func_error(
                 self.cfg.input.clone(),
-                pyloc_to_ergloc(range),
+                pyloc_to_ergloc(func_def.range),
-                self.namespace.join("."),
+                self.cur_namespace(),
                 &name,
             );
             self.errs.push(err);
             Expr::Dummy(Dummy::new(None, vec![]))
         } else {
-            let loc = range.start;
+            let loc = func_def.range.start;
-            let decos = decorator_list
+            let decos = func_def
+                .decorator_list
                 .into_iter()
                 .map(|ex| Decorator(self.convert_expr(ex)))
                 .collect::<HashSet<_>>();
@@ -1613,14 +1727,9 @@ impl ASTConverter {
                 );
                 TypeSpecWithOp::new(colon, t_spec, self.convert_expr(ret))
             });
-            let sig = Signature::Subr(SubrSignature::new(
+            let bounds = self.get_type_bounds(func_def.type_params);
-                decos,
+            let sig = Signature::Subr(SubrSignature::new(decos, ident, bounds, params, return_t));
-                ident,
+            let block = self.convert_block(func_def.body, BlockKind::Function);
-                TypeBoundSpecs::empty(),
-                params,
-                return_t,
-            ));
-            let block = self.convert_block(body, BlockKind::Function);
             let body = DefBody::new(EQUAL, block, DefId(0));
             let def = Def::new(sig, body);
             self.pop();
@@ -1642,19 +1751,16 @@ impl ASTConverter {
     /// ```erg
     /// Foo = Inherit Bar
     /// ```
-    fn convert_classdef(
+    fn convert_classdef(&mut self, class_def: py_ast::StmtClassDef) -> Expr {
-        &mut self,
+        let loc = class_def.location();
-        name: String,
+        let name = class_def.name.to_string();
-        body: Vec<py_ast::Stmt>,
+        let _decos = class_def
-        bases: Vec<py_ast::Expr>,
+            .decorator_list
-        decorator_list: Vec<py_ast::Expr>,
-        loc: PyLocation,
-    ) -> Expr {
-        let _decos = decorator_list
             .into_iter()
             .map(|deco| self.convert_expr(deco))
             .collect::<Vec<_>>();
-        let mut bases = bases
+        let mut bases = class_def
+            .bases
             .into_iter()
             .map(|base| self.convert_expr(base))
             .collect::<Vec<_>>();
@@ -1667,7 +1773,7 @@ impl ASTConverter {
         let ident = self.convert_ident(name, class_name_loc);
         let sig = Signature::Var(VarSignature::new(VarPattern::Ident(ident.clone()), None));
         self.grow(ident.inspect().to_string());
-        let (base_type, methods) = self.extract_method_list(ident, body, inherit);
+        let (base_type, methods) = self.extract_method_list(ident, class_def.body, inherit);
         let classdef = if inherit {
             // TODO: multiple inheritance
             let pos_args = vec![PosArg::new(bases.remove(0))];
@@ -1763,6 +1869,20 @@ impl ASTConverter {
                     match lhs {
                         py_ast::Expr::Name(name) => {
                             let expr = self.convert_expr(*assign.value);
+                            if let Expr::Call(call) = &expr {
+                                if let Some("TypeVar") = call.obj.get_name().map(|s| &s[..]) {
+                                    let arg = if let Some(Expr::Literal(lit)) =
+                                        call.args.get_left_or_key("arg")
+                                    {
+                                        lit.token.content.trim_matches('\"').to_string()
+                                    } else {
+                                        name.id.to_string()
+                                    };
+                                    let bound = call.args.nth_or_key(1, "bound").cloned();
+                                    let info = TypeVarInfo::new(arg, bound);
+                                    self.define_type_var(name.id.to_string(), info);
+                                }
+                            }
                             let can_shadow = self.register_name_info(&name.id, NameKind::Variable);
                             let ident = self.convert_ident(name.id.to_string(), name.location());
                             if can_shadow.is_yes() {
@@ -1918,24 +2038,8 @@ impl ASTConverter {
                     }
                 }
             }
-            py_ast::Stmt::FunctionDef(func_def) => self.convert_funcdef(
+            py_ast::Stmt::FunctionDef(func_def) => self.convert_funcdef(func_def),
-                func_def.name.to_string(),
+            py_ast::Stmt::ClassDef(class_def) => self.convert_classdef(class_def),
-                *func_def.args,
-                func_def.body,
-                func_def.decorator_list,
-                func_def.returns.map(|x| *x),
-                func_def.range,
-            ),
-            py_ast::Stmt::ClassDef(class_def) => {
-                let class_loc = class_def.location();
-                self.convert_classdef(
-                    class_def.name.to_string(),
-                    class_def.body,
-                    class_def.bases,
-                    class_def.decorator_list,
-                    class_loc,
-                )
-            }
             py_ast::Stmt::For(for_) => {
                 let loc = for_.location();
                 let iter = self.convert_expr(*for_.iter);
@@ -1991,7 +2095,7 @@ impl ASTConverter {
                     value
                 } else {
                     let func_acc = Expr::Accessor(Accessor::Ident(
-                        self.convert_ident(self.namespace.last().unwrap().clone(), loc),
+                        self.convert_ident(self.cur_name().to_string(), loc),
                     ));
                     let return_acc = self.convert_ident("return".to_string(), loc);
                     let return_acc = Expr::Accessor(Accessor::attr(func_acc, return_acc));

crates/py2erg/gen_decl.rs

@@ -71,7 +71,10 @@ impl DeclFileGenerator {
     // e.g. `x: foo.Bar` => `foo = pyimport "foo"; x: foo.Bar`
     fn prepare_using_type(&mut self, typ: &Type) {
         let namespace = Str::rc(typ.namespace().split('.').next().unwrap());
-        if namespace != self.namespace && !namespace.is_empty() && self.imported.insert(namespace.clone()) {
+        if namespace != self.namespace
+            && !namespace.is_empty()
+            && self.imported.insert(namespace.clone())
+        {
             self.code += &format!("{namespace} = pyimport \"{namespace}\"\n");
         }
     }

tests/test.rs

@@ -137,6 +137,11 @@ fn exec_shadowing() -> Result<(), String> {
     expect("tests/shadowing.py", 0, 3)
 }
 
+#[test]
+fn exec_typevar() -> Result<(), String> {
+    expect("tests/typevar.py", 0, 2)
+}
+
 #[test]
 fn exec_widening() -> Result<(), String> {
     expect("tests/widening.py", 0, 1)

tests/typevar.py

@@ -0,0 +1,25 @@
+from typing import TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U", int)
+def id(x: T) -> T:
+    return x
+
+def id_int(x: U) -> U:
+    return x
+
+_ = id(1) + 1 # OK
+_ = id("a") + "b" # OK
+_ = id_int(1) # OK
+_ = id_int("a") # ERR
+
+def id2[T](x: T) -> T:
+    return x
+
+def id_int2[T: int](x: T) -> T:
+    return x
+
+_ = id2(1) + 1  # OK
+_ = id2("a") + "b"  # OK
+_ = id_int2(1) # OK
+_ = id_int2("a") # ERR