SONARPY-1485 Support type parameter syntax for functions (#1602)

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/BaseTreeVisitor.java

@@ -133,6 +133,7 @@ public void visitContinueStatement(ContinueStatement pyContinueStatementTree) {
   public void visitFunctionDef(FunctionDef pyFunctionDefTree) {
     scan(pyFunctionDefTree.decorators());
     scan(pyFunctionDefTree.name());
+    scan(pyFunctionDefTree.typeParams());
     scan(pyFunctionDefTree.parameters());
     scan(pyFunctionDefTree.returnTypeAnnotation());
     scan(pyFunctionDefTree.body());
@@ -561,4 +562,15 @@ public void visitKeywordPattern(KeywordPattern keywordPattern) {
   public void visitValuePattern(ValuePattern valuePattern) {
     scan(valuePattern.qualifiedExpression());
   }
+
+  @Override
+  public void visitTypeParams(TypeParams typeParams) {
+    scan(typeParams.typeParamsList());
+  }
+
+  @Override
+  public void visitTypeParam(TypeParam typeParam) {
+    scan(typeParam.name());
+    scan(typeParam.typeAnnotation());
+  }
 }

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/FunctionDef.java

@@ -43,6 +43,9 @@ public interface FunctionDef extends Statement, FunctionLike {
 
   Name name();
 
+  @CheckForNull
+  TypeParams typeParams();
+
   Token leftPar();
 
   Token rightPar();

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/Tree.java

@@ -179,7 +179,10 @@ enum Kind {
 
     VARIABLE_TYPE_ANNOTATION(TypeAnnotation.class),
     PARAMETER_TYPE_ANNOTATION(TypeAnnotation.class),
+    TYPE_PARAM_TYPE_ANNOTATION(TypeAnnotation.class),
     RETURN_TYPE_ANNOTATION(TypeAnnotation.class),
+    TYPE_PARAMS(TypeParams.class),
+    TYPE_PARAM(TypeParam.class),
 
     PARAMETER_LIST(ParameterList.class),
     VALUE_PATTERN(ValuePattern.class),

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/TreeVisitor.java

@@ -206,4 +206,8 @@ public interface TreeVisitor {
   void visitKeywordPattern(KeywordPattern keywordPattern);
 
   void visitValuePattern(ValuePattern valuePattern);
+
+  void visitTypeParams(TypeParams typeParams);
+
+  void visitTypeParam(TypeParam typeParam);
 }

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/TypeParam.java

@@ -0,0 +1,34 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2023 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.plugins.python.api.tree;
+
+import javax.annotation.CheckForNull;
+
+public interface TypeParam extends Tree {
+
+  @CheckForNull
+  Token starToken();
+
+  Name name();
+
+  @CheckForNull
+  TypeAnnotation typeAnnotation();
+
+}

python-frontend/src/main/java/org/sonar/plugins/python/api/tree/TypeParams.java

@@ -0,0 +1,28 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2023 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.plugins.python.api.tree;
+
+import java.util.List;
+
+public interface TypeParams extends Tree {
+  Token leftBracket();
+  List<TypeParam> typeParamsList();
+  Token rightBracket();
+}

python-frontend/src/main/java/org/sonar/python/api/PythonGrammar.java

@@ -180,6 +180,8 @@ public enum PythonGrammar implements GrammarRuleKey {
   ATTR,
   FUNCNAME,
   FUN_RETURN_ANNOTATION,
+  TYPE_PARAMS,
+  TYPE_PARAM,
 
   CLASSDEF,
   CLASSNAME,

python-frontend/src/main/java/org/sonar/python/api/PythonGrammarBuilder.java

@@ -385,8 +385,18 @@ protected void compoundStatements(LexerfulGrammarBuilder b) {
     b.rule(WILDCARD_PATTERN).is("_");
     b.rule(GROUP_PATTERN).is("(", PATTERN, ")");
 
-    b.rule(FUNCDEF).is(b.optional(DECORATORS), b.optional(ASYNC), "def", FUNCNAME, "(", b.optional(TYPEDARGSLIST), ")", b.optional(FUN_RETURN_ANNOTATION), ":", SUITE);
+    b.rule(FUNCDEF).is(b.optional(DECORATORS),
+      b.optional(ASYNC),
+      "def",
+      FUNCNAME,
+      b.optional(TYPE_PARAMS),
+      "(", b.optional(TYPEDARGSLIST), ")",
+      b.optional(FUN_RETURN_ANNOTATION),
+      ":",
+      SUITE);
+
     b.rule(FUNCNAME).is(NAME);
+    b.rule(TYPE_PARAMS).is("[", TYPEDARGSLIST, "]");
     b.rule(FUN_RETURN_ANNOTATION).is("-", ">", TEST);
 
     b.rule(DECORATORS).is(b.oneOrMore(DECORATOR));

python-frontend/src/main/java/org/sonar/python/tree/FunctionDefImpl.java

@@ -40,13 +40,15 @@
 import org.sonar.plugins.python.api.tree.TreeVisitor;
 import org.sonar.plugins.python.api.tree.TypeAnnotation;
 import org.sonar.plugins.python.api.symbols.Symbol;
+import org.sonar.plugins.python.api.tree.TypeParams;
 
 public class FunctionDefImpl extends PyTree implements FunctionDef {
 
   private final List<Decorator> decorators;
   private final Token asyncKeyword;
   private final Token defKeyword;
   private final Name name;
+  private final TypeParams typeParams;
   private final Token leftPar;
   private final ParameterList parameters;
   private final Token rightPar;
@@ -62,14 +64,15 @@ public class FunctionDefImpl extends PyTree implements FunctionDef {
   private Set<Symbol> symbols = new HashSet<>();
   private FunctionSymbol functionSymbol;
 
-  public FunctionDefImpl(List<Decorator> decorators, @Nullable Token asyncKeyword, Token defKeyword, Name name,
+  public FunctionDefImpl(List<Decorator> decorators, @Nullable Token asyncKeyword, Token defKeyword, Name name, TypeParams typeParams,
                          Token leftPar, @Nullable ParameterList parameters, Token rightPar, @Nullable TypeAnnotation returnType,
                          Token colon, @Nullable Token newLine, @Nullable Token indent, StatementList body, @Nullable Token dedent,
                          boolean isMethodDefinition, @Nullable StringLiteral docstring) {
     this.decorators = decorators;
     this.asyncKeyword = asyncKeyword;
     this.defKeyword = defKeyword;
     this.name = name;
+    this.typeParams = typeParams;
     this.leftPar = leftPar;
     this.parameters = parameters;
     this.rightPar = rightPar;
@@ -104,6 +107,12 @@ public Name name() {
     return name;
   }
 
+  @CheckForNull
+  @Override
+  public TypeParams typeParams() {
+    return typeParams;
+  }
+
   @Override
   public Token leftPar() {
     return leftPar;
@@ -168,7 +177,7 @@ public void accept(TreeVisitor visitor) {
 
   @Override
   public List<Tree> computeChildren() {
-    return Stream.of(decorators, Arrays.asList(asyncKeyword, defKeyword, name, leftPar, parameters, rightPar, returnType, colon, newLine, indent, body, dedent))
+    return Stream.of(decorators, Arrays.asList(asyncKeyword, defKeyword, name, typeParams, leftPar, parameters, rightPar, returnType, colon, newLine, indent, body, dedent))
       .flatMap(List::stream).filter(Objects::nonNull).collect(Collectors.toList());
   }
 

python-frontend/src/main/java/org/sonar/python/tree/PythonTreeMaker.java

@@ -23,9 +23,11 @@
 import com.sonar.sslr.api.GenericTokenType;
 import com.sonar.sslr.api.RecognitionException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.CheckForNull;
 import javax.annotation.Nullable;
@@ -98,6 +100,8 @@
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.plugins.python.api.tree.TryStatement;
 import org.sonar.plugins.python.api.tree.TypeAnnotation;
+import org.sonar.plugins.python.api.tree.TypeParam;
+import org.sonar.plugins.python.api.tree.TypeParams;
 import org.sonar.plugins.python.api.tree.WithItem;
 import org.sonar.plugins.python.api.tree.WithStatement;
 import org.sonar.plugins.python.api.tree.YieldExpression;
@@ -554,14 +558,10 @@ public FunctionDef funcDefStatement(AstNode astNode) {
         .collect(Collectors.toList());
     }
     Name name = name(astNode.getFirstChild(PythonGrammar.FUNCNAME).getFirstChild(PythonGrammar.NAME));
-    ParameterList parameterList = null;
-    AstNode typedArgListNode = astNode.getFirstChild(PythonGrammar.TYPEDARGSLIST);
-    if (typedArgListNode != null) {
-      List<AnyParameter> arguments = typedArgListNode.getChildren(PythonGrammar.TFPDEF, PythonPunctuator.MUL, PythonPunctuator.DIV).stream()
-        .map(this::parameter).filter(Objects::nonNull).collect(Collectors.toList());
-      List<Token> commas = punctuators(typedArgListNode, PythonPunctuator.COMMA);
-      parameterList = new ParameterListImpl(arguments, commas);
-    }
+
+    var typeParams = typeParams(astNode);
+
+    var parameterList = parameterList(astNode);
 
     AstNode suite = astNode.getFirstChild(PythonGrammar.SUITE);
     StatementList body = getStatementListFromSuite(suite);
@@ -582,11 +582,70 @@ public FunctionDef funcDefStatement(AstNode astNode) {
     }
 
     Token colon = toPyToken(astNode.getFirstChild(PythonPunctuator.COLON).getToken());
-    return new FunctionDefImpl(decorators, asyncToken, toPyToken(defNode.getToken()), name, lPar, parameterList, rPar,
+    return new FunctionDefImpl(decorators, asyncToken, toPyToken(defNode.getToken()), name, typeParams, lPar, parameterList, rPar,
       returnType, colon, suiteNewLine(suite), suiteIndent(suite), body, suiteDedent(suite),
       isMethodDefinition(astNode), DocstringExtractor.extractDocstring(body));
   }
 
+  private ParameterList parameterList(AstNode parent) {
+    return Optional.of(parent)
+      .map(n -> n.getFirstChild(PythonGrammar.TYPEDARGSLIST))
+      .map(n -> {
+        List<AnyParameter> arguments = n.getChildren(PythonGrammar.TFPDEF, PythonPunctuator.MUL, PythonPunctuator.DIV).stream()
+          .map(this::parameter).filter(Objects::nonNull).collect(Collectors.toList());
+        List<Token> commas = punctuators(n, PythonPunctuator.COMMA);
+        return new ParameterListImpl(arguments, commas);
+      }).orElse(null);
+  }
+
+  private TypeParams typeParams(AstNode parent) {
+    return Optional.of(parent)
+      .map(n -> n.getFirstChild(PythonGrammar.TYPE_PARAMS))
+      .map(n -> {
+        var lBracket = toPyToken(n.getFirstChild(PythonPunctuator.LBRACKET).getToken());
+
+        var parameters = Optional.of(n.getFirstChild(PythonGrammar.TYPEDARGSLIST))
+          .map(argList -> argList.getChildren(PythonGrammar.TFPDEF))
+          .stream()
+          .flatMap(Collection::stream)
+          .map(this::typeParam)
+          .collect(Collectors.toList());
+
+        var commas = Optional.of(n.getFirstChild(PythonGrammar.TYPEDARGSLIST))
+          .map(argList -> punctuators(argList, PythonPunctuator.COMMA))
+          .stream()
+          .flatMap(Collection::stream)
+          .collect(Collectors.toList());
+
+        var rBracket = toPyToken(n.getFirstChild(PythonPunctuator.RBRACKET).getToken());
+
+        return new TypeParamsImpl(lBracket, parameters, commas, rBracket);
+      }).orElse(null);
+  }
+
+  private TypeParam typeParam(AstNode parameter) {
+    var starOrStarStar = Optional.of(parameter)
+      .map(AstNode::getPreviousSibling)
+      .filter(ps -> ps.is(PythonPunctuator.MUL, PythonPunctuator.MUL_MUL))
+      .map(ps -> toPyToken(ps.getToken()))
+      .orElse(null);
+
+    Name name = name(parameter.getFirstChild(PythonGrammar.NAME));
+
+    var typeAnnotation = Optional.of(parameter)
+      .filter(p -> Objects.nonNull(p.getFirstChild(PythonPunctuator.COLON)))
+      .filter(p -> Objects.nonNull(p.getFirstChild(PythonGrammar.TEST)))
+      .map(p -> {
+        var colonNode = parameter.getFirstChild(PythonPunctuator.COLON);
+        var testNode = parameter.getFirstChild(PythonGrammar.TEST);
+        var colonToken = toPyToken(colonNode.getToken());
+        var testExpression = expression(testNode);
+        return new TypeAnnotationImpl(colonToken, testExpression, Tree.Kind.TYPE_PARAM_TYPE_ANNOTATION);
+      }).orElse(null);
+
+    return new TypeParamImpl(starOrStarStar, name, typeAnnotation);
+  }
+
   private Decorator decorator(AstNode astNode) {
     Token atToken = toPyToken(astNode.getFirstChild(PythonPunctuator.AT).getToken());
     Expression expression = expression(astNode.getFirstChild(PythonGrammar.NAMED_EXPR_TEST));

python-frontend/src/main/java/org/sonar/python/tree/TypeParamImpl.java

@@ -0,0 +1,78 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2023 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python.tree;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import javax.annotation.CheckForNull;
+import javax.annotation.Nullable;
+import org.sonar.plugins.python.api.tree.Name;
+import org.sonar.plugins.python.api.tree.Token;
+import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.plugins.python.api.tree.TreeVisitor;
+import org.sonar.plugins.python.api.tree.TypeAnnotation;
+import org.sonar.plugins.python.api.tree.TypeParam;
+
+public class TypeParamImpl extends PyTree implements TypeParam {
+
+  private final Token starToken;
+  private final Name name;
+  private final TypeAnnotation annotation;
+
+  public TypeParamImpl(@Nullable Token starToken, Name name, @Nullable TypeAnnotation annotation) {
+    this.starToken = starToken;
+    this.name = name;
+    this.annotation = annotation;
+  }
+
+  @Override
+  public void accept(TreeVisitor visitor) {
+    visitor.visitTypeParam(this);
+  }
+
+  @Override
+  public Kind getKind() {
+    return Kind.TYPE_PARAM;
+  }
+
+  @CheckForNull
+  @Override
+  public Token starToken() {
+    return starToken;
+  }
+
+  @Override
+  public Name name() {
+    return name;
+  }
+
+  @CheckForNull
+  @Override
+  public TypeAnnotation typeAnnotation() {
+    return annotation;
+  }
+
+  @Override
+  List<Tree> computeChildren() {
+    return Stream.of(starToken, name, annotation).filter(Objects::nonNull).collect(Collectors.toList());
+  }
+}