From 4fc97d088fa25adca5ad2a2fcf42f08f49b39e74 Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Mon, 18 Dec 2017 18:24:22 +0300 Subject: [PATCH] Correct types for initvars in `__init__` (PY-27398) --- .../stdlib/PyDataclassesTypeProvider.kt | 15 +++++++++++++-- .../InitializingDataclass/a.py | 11 +++++++++++ .../InitializingDataclass/a.py | 11 ++++++++++- .../InitializingDataclass/dataclasses.py | 8 ++++++++ .../testData/paramInfo/InitializingDataclass/a.py | 10 +++++++++- .../InitializingDataclass/dataclasses.py | 8 ++++++++ .../com/jetbrains/python/PyParameterInfoTest.java | 3 ++- 7 files changed, 61 insertions(+), 5 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt index 1738568a67e0..94c225acec7d 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt @@ -4,7 +4,10 @@ package com.jetbrains.python.codeInsight.stdlib import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider -import com.jetbrains.python.psi.* +import com.jetbrains.python.psi.PyClass +import com.jetbrains.python.psi.PyReferenceExpression +import com.jetbrains.python.psi.PyTargetExpression +import com.jetbrains.python.psi.PyUtil import com.jetbrains.python.psi.impl.PyCallExpressionNavigator import com.jetbrains.python.psi.resolve.PyResolveContext import com.jetbrains.python.psi.types.* @@ -41,7 +44,7 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() { val annotation = element.annotation if (annotation != null && !PyTypingTypeProvider.isClassVarAnnotation(annotation, context)) { - parameters.add(PyCallableParameterImpl.nonPsi(element.name, context.getType(element), element.findAssignedValue())) + parameters.add(PyCallableParameterImpl.nonPsi(element.name, getTypeForParameter(element, context), element.findAssignedValue())) } } @@ -50,4 +53,12 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() { return PyCallableTypeImpl(parameters, context.getType(cls)) } + + private fun getTypeForParameter(element: PyTargetExpression, context: TypeEvalContext): PyType? { + val type = context.getType(element) + if (type is PyCollectionType && type is PyClassType && type.classQName == "dataclasses.InitVar") { + return type.elementTypes.firstOrNull() + } + return type + } } \ No newline at end of file diff --git a/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py b/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py index 6fbe96236e7c..56e4c02cf63d 100644 --- a/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py +++ b/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py @@ -74,3 +74,14 @@ class C2: C2() C2(1) C2(1, 2) + + +@dataclasses.dataclass +class D1: + a: dataclasses.InitVar[int] + b: int + +D1() +D1(1) +D1(1, 2) +D1(1, 2, 3) diff --git a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py index 6f06d95973cf..a54506c4f65d 100644 --- a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py +++ b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py @@ -73,4 +73,13 @@ class C2: b: int C2(1) -C2("1") \ No newline at end of file +C2("1") + + +@dataclasses.dataclass +class D1: + a: dataclasses.InitVar[int] + b: int + +D1(1, 2) +D1("1", "2") \ No newline at end of file diff --git a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py index 3a2b30d36707..e2c6ed38f1c6 100644 --- a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py +++ b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py @@ -1,3 +1,11 @@ +class _InitVarMeta(type): + def __getitem__(self, params): + return self + +class InitVar(metaclass=_InitVarMeta): + pass + + def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False, hash=None, frozen=False): pass \ No newline at end of file diff --git a/python/testData/paramInfo/InitializingDataclass/a.py b/python/testData/paramInfo/InitializingDataclass/a.py index 792a5b7b67b6..fad9eaf9905e 100644 --- a/python/testData/paramInfo/InitializingDataclass/a.py +++ b/python/testData/paramInfo/InitializingDataclass/a.py @@ -55,4 +55,12 @@ class C2: a: typing.ClassVar b: int -C2() \ No newline at end of file +C2() + + +@dataclasses.dataclass +class D1: + a: dataclasses.InitVar[int] + b: int + +D1() \ No newline at end of file diff --git a/python/testData/paramInfo/InitializingDataclass/dataclasses.py b/python/testData/paramInfo/InitializingDataclass/dataclasses.py index 3a2b30d36707..e2c6ed38f1c6 100644 --- a/python/testData/paramInfo/InitializingDataclass/dataclasses.py +++ b/python/testData/paramInfo/InitializingDataclass/dataclasses.py @@ -1,3 +1,11 @@ +class _InitVarMeta(type): + def __getitem__(self, params): + return self + +class InitVar(metaclass=_InitVarMeta): + pass + + def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False, hash=None, frozen=False): pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java b/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java index 226263242df5..e29412192fb1 100644 --- a/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java +++ b/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java @@ -712,7 +712,7 @@ public class PyParameterInfoTest extends LightMarkedTestCase { runWithLanguageLevel( LanguageLevel.PYTHON37, () -> { - final Map marks = loadMultiFileTest(6); + final Map marks = loadMultiFileTest(7); feignCtrlP(marks.get("").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "}); feignCtrlP(marks.get("").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "}); @@ -726,6 +726,7 @@ public class PyParameterInfoTest extends LightMarkedTestCase { feignCtrlP(marks.get("").getTextOffset()).check("self: B2, x: int", new String[]{"x: int"}, new String[]{"self: B2, "}); feignCtrlP(marks.get("").getTextOffset()).check("b: int", new String[]{"b: int"}); feignCtrlP(marks.get("").getTextOffset()).check("b: int", new String[]{"b: int"}); + feignCtrlP(marks.get("").getTextOffset()).check("a: int, b: int", new String[]{"a: int, "}); } ); }