From 5e214c4b4af2017f4f7d43ba6e7a16c070854ff9 Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Mon, 21 Nov 2016 19:52:05 +0300 Subject: [PATCH] PY-19723 Fixed: Type hinting of arbitrary argument lists and default argument values Update PyTypingTypeProvider to correctly handle `args: T1` and `kwargs: T2`, and wrap their types to `Tuple[T1, ...]` and `Dict[str, T2]`. --- .../codeInsight/PyTypingTypeProvider.java | 87 ++++++++++++------- .../python/psi/types/PyTypeUtil.java | 22 ++++- .../com/jetbrains/python/PyTypingTest.java | 29 ++++++- 3 files changed, 106 insertions(+), 32 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java index fd424378e7e6..dcba9d1781d4 100644 --- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java @@ -79,22 +79,14 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { @Nullable public Ref getParameterType(@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) { - final PyAnnotation annotation = param.getAnnotation(); - if (annotation != null) { - // XXX: Requires switching from stub to AST - final PyExpression value = annotation.getValue(); - if (value != null) { - final PyType type = getType(value, new Context(context)); - if (type != null) { - final PyType optionalType = getOptionalTypeFromDefaultNone(param, type, context); - return Ref.create(optionalType != null ? optionalType : type); - } - } + final Ref typeFromAnnotation = getParameterTypeFromAnnotation(param, context); + if (typeFromAnnotation != null) { + return typeFromAnnotation; } - final String paramComment = param.getTypeCommentAnnotation(); - if (paramComment != null) { - return Ref.create(getStringBasedType(paramComment, param, new Context(context))); + final Ref typeFromTypeComment = getParameterTypeFromTypeComment(param, context); + if (typeFromTypeComment != null) { + return typeFromTypeComment; } final String comment = func.getTypeCommentAnnotation(); @@ -123,6 +115,57 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { return null; } + @Nullable + private static Ref getParameterTypeFromAnnotation(@NotNull PyNamedParameter parameter, @NotNull TypeEvalContext context) { + final PyType annotationValueType = Optional + .ofNullable(parameter.getAnnotation()) + .map(PyAnnotation::getValue) // XXX: Requires switching from stub to AST + .map(value -> getType(value, new Context(context))) + .orElse(null); + + if (annotationValueType != null) { + if (parameter.isPositionalContainer()) { + return Ref.create(PyTypeUtil.toPositionalContainerType(parameter, annotationValueType)); + } + + if (parameter.isKeywordContainer()) { + return Ref.create(PyTypeUtil.toKeywordContainerType(parameter, annotationValueType)); + } + + final PyType result = Optional + .ofNullable(parameter.getDefaultValue()) + .map(context::getType) + .filter(PyNoneType.class::isInstance) + .map(noneType -> PyUnionType.union(annotationValueType, noneType)) + .orElse(annotationValueType); + + return Ref.create(result); + } + + return null; + } + + @Nullable + private static Ref getParameterTypeFromTypeComment(@NotNull PyNamedParameter parameter, @NotNull TypeEvalContext context) { + final String typeComment = parameter.getTypeCommentAnnotation(); + + if (typeComment != null) { + final PyType type = getStringBasedType(typeComment, parameter, new Context(context)); + + if (parameter.isPositionalContainer()) { + return Ref.create(PyTypeUtil.toPositionalContainerType(parameter, type)); + } + + if (parameter.isKeywordContainer()) { + return Ref.create(PyTypeUtil.toKeywordContainerType(parameter, type)); + } + + return Ref.create(type); + } + + return null; + } + private static boolean omitFirstParamInTypeComment(@NotNull PyFunction func) { return func.getContainingClass() != null && func.getModifier() != PyFunction.Modifier.STATICMETHOD; } @@ -204,7 +247,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } /** - * Checks that text of a comment starts with the "type:" prefix and returns trimmed part afterwards. This trailing part is supposed to + * Checks that text of a comment starts with the "type:" prefix and returns trimmed part afterwards. This trailing part is supposed to * contain type annotation in PEP 484 compatible format, that can be parsed with either {@link PyTypeParser#parse(PsiElement, String)} * or {@link PyTypeParser#parsePep484FunctionTypeComment(PsiElement, String)}. */ @@ -221,20 +264,6 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { return type instanceof PyClassType && "typing.Any".equals(((PyClassType)type).getPyClass().getQualifiedName()); } - @Nullable - private static PyType getOptionalTypeFromDefaultNone(@NotNull PyNamedParameter param, - @NotNull PyType type, - @NotNull TypeEvalContext context) { - final PyExpression defaultValue = param.getDefaultValue(); - if (defaultValue != null) { - final PyType defaultType = context.getType(defaultValue); - if (defaultType instanceof PyNoneType) { - return PyUnionType.union(type, defaultType); - } - } - return null; - } - @Nullable private static PyType getGenericConstructorType(@NotNull PyFunction function, @NotNull Context context) { if (PyUtil.isInit(function)) { diff --git a/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java b/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java index 3702a45f9368..1f4714a79959 100644 --- a/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java +++ b/python/src/com/jetbrains/python/psi/types/PyTypeUtil.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2015 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,12 +18,14 @@ package com.jetbrains.python.psi.types; import com.intellij.openapi.util.Key; import com.intellij.openapi.util.UserDataHolder; import com.intellij.psi.PsiElement; -import com.intellij.util.Processor; +import com.jetbrains.python.psi.impl.PyBuiltinCache; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; +import java.util.Optional; /** * Tools and wrappers around {@link PyType} inheritors @@ -81,4 +83,20 @@ public final class PyTypeUtil { } return null; } + + @Nullable + public static PyTupleType toPositionalContainerType(@NotNull PsiElement anchor, @Nullable PyType elementType) { + return PyTupleType.createHomogeneous(anchor, elementType); + } + + @Nullable + public static PyCollectionType toKeywordContainerType(@NotNull PsiElement anchor, @Nullable PyType valueType) { + final PyBuiltinCache builtinCache = PyBuiltinCache.getInstance(anchor); + + return Optional + .ofNullable(builtinCache.getDictType()) + .map(PyClassType::getPyClass) + .map(dictClass -> new PyCollectionTypeImpl(dictClass, false, Arrays.asList(builtinCache.getStrType(), valueType))) + .orElse(null); + } } diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index c216449ca0ac..947f424c896e 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -31,7 +31,6 @@ import com.jetbrains.python.psi.LanguageLevel; import com.jetbrains.python.psi.PyExpression; import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.TypeEvalContext; -import org.intellij.lang.annotations.Language; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -736,6 +735,34 @@ public class PyTypingTest extends PyTestCase { "expr = (w, x, y, z)\n"); } + // PY-19723 + public void testAnnotatedPositionalArgs() { + doTest("Tuple[str, ...]", + "def foo(*args: str):\n" + + " expr = args\n"); + } + + // PY-19723 + public void testAnnotatedKeywordArgs() { + doTest("Dict[str, int]", + "def foo(**kwargs: int):\n" + + " expr = kwargs\n"); + } + + // PY-19723 + public void testTypeCommentedPositionalArgs() { + doTest("Tuple[str, ...]", + "def foo(*args # type: str\n):\n" + + " expr = args\n"); + } + + // PY-19723 + public void testTypeCommentedKeywordArgs() { + doTest("Dict[str, int]", + "def foo(**kwargs # type: int\n):\n" + + " expr = kwargs\n"); + } + private void doTestNoInjectedText(@NotNull String text) { myFixture.configureByText(PythonFileType.INSTANCE, text); final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());