Support type parameters

marijnh · marijnh · commit 05886a2daf6e · 2023-10-20T11:27:41.000+02:00
FIX: Implement PEP 695 type parameter syntax. Closes #18 Closes #19
diff --git a/src/python.grammar b/src/python.grammar
@@ -22,14 +22,21 @@
 Decorator { At dottedName ArgList? newline }
 
 FunctionDefinition {
-  kw<"async">? kw<"def"> VariableName
+  kw<"async">? kw<"def"> VariableName TypeParamList?
   ParamList
   TypeDef { "->" test }?
   Body
 }
 
 ParamList { "(" commaSep<param>? ")" }
 
+TypeParamList { "[" commaSep<TypeParam> "]" }
+
+TypeParam {
+  ("*" | "**") VariableName |
+  VariableName TypeDef?
+}
+
 MatchStatement {
   skw<"match"> "*"? expression MatchBody { ":" newline indent MatchClause+ (dedent | eof) }
 }
@@ -65,7 +72,7 @@ SequencePattern { "[" commaSep<pattern> "]" | "(" commaSep<pattern> ")" }
 
 MappingPattern { "{" commaSep<"**" pattern | (VariableName | LiteralPattern) ":" pattern> "}" }
 
-ClassDefinition { kw<"class"> VariableName ArgList? Body }
+ClassDefinition { kw<"class"> VariableName TypeParamList? ArgList? Body }
 
 param { VariableName TypeDef? (AssignOp{"="} test)? | "*" VariableName? | "**" VariableName | "/" }
 
@@ -92,7 +99,8 @@ smallStatement {
   RaiseStatement { kw<"raise"> (test (kw<"from"> test | ("," test ("," test)?))?)? } |
   ImportStatement |
   ScopeStatement { (kw<"global"> | kw<"nonlocal">) commaSep<VariableName> } |
-  AssertStatement { kw<"assert"> commaSep<test> }
+  AssertStatement { kw<"assert"> commaSep<test> } |
+  TypeDefinition { skw<"type"> VariableName TypeParamList? "=" test }
 }
 
 expressions { commaSep<"*" expression | test> }
diff --git a/src/tokens.js b/src/tokens.js
@@ -2,7 +2,7 @@ import {ExternalTokenizer, ContextTracker} from "@lezer/lr"
 import {
   newline as newlineToken, eof, newlineBracketed, blankLineStart, indent, dedent, printKeyword,
   ParenthesizedExpression, TupleExpression, ComprehensionExpression,
-  PatternArgList, SequencePattern, MappingPattern, FormatString,
+  PatternArgList, SequencePattern, MappingPattern, FormatString, TypeParamList,
   ArrayExpression, ArrayComprehensionExpression, ArgList, ParamList, importList, subscript,
   DictionaryExpression, DictionaryComprehensionExpression, SetExpression, SetComprehensionExpression,
   formatString1Content, formatString1Brace, formatString1End,
@@ -20,7 +20,7 @@ const bracketed = new Set([
   ArrayExpression, ArrayComprehensionExpression, subscript,
   SetExpression, SetComprehensionExpression, FormatString,
   DictionaryExpression, DictionaryComprehensionExpression,
-  SequencePattern, MappingPattern, PatternArgList
+  SequencePattern, MappingPattern, PatternArgList, TypeParamList
 ])
 
 function isLineBreak(ch) {
diff --git a/test/statement.txt b/test/statement.txt
@@ -383,3 +383,32 @@ Script(MatchStatement(match,VariableName,MatchBody(
       LogicOp,
       SequencePattern(CapturePattern(VariableName),CapturePattern(VariableName))),
     Body(PassStatement(pass))))))
+
+# Type Params
+
+class ClassA[T: str]:
+  def method1(self) -> T:
+    pass
+
+def func[**T](a: T, b: T) -> T:
+  pass
+
+==>
+
+Script(
+  ClassDefinition(class,VariableName,
+    TypeParamList(TypeParam(VariableName,TypeDef(":",VariableName))),
+    Body(FunctionDefinition(def,VariableName,ParamList(VariableName),TypeDef(VariableName),
+      Body(PassStatement(pass))))),
+  FunctionDefinition(def,VariableName,TypeParamList(TypeParam(VariableName)),
+    ParamList(VariableName,TypeDef(VariableName),VariableName,TypeDef(VariableName)),
+    TypeDef(VariableName),
+    Body(PassStatement(pass))))
+
+# Type Definition
+
+type Point = tuple[float, float]
+
+==>
+
+Script(TypeDefinition(type,VariableName,MemberExpression(VariableName,VariableName,",",VariableName)))
