From b6982b44259a8ff7782bc0411a844e3063291fe5 Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Thu, 14 Dec 2017 15:34:23 +0300 Subject: [PATCH] Support `@dataclass(init=True/False)` (PY-27398) --- .../codeInsight/stdlib/PyDataclassUtil.kt | 115 ++++++++++++++++++ .../stdlib/PyDataclassesTypeProvider.kt | 3 +- .../InitializingDataclass/a.py | 44 ++++++- .../InitializingDataclass/a.py | 46 ++++++- .../paramInfo/InitializingDataclass/a.py | 34 +++++- .../jetbrains/python/PyParameterInfoTest.java | 11 +- 6 files changed, 248 insertions(+), 5 deletions(-) create mode 100644 python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt new file mode 100644 index 000000000000..3222fc4b6e03 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt @@ -0,0 +1,115 @@ +/* + * Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. + */ +package com.jetbrains.python.codeInsight.stdlib + +import com.jetbrains.python.psi.* +import com.jetbrains.python.psi.impl.PyEvaluator +import com.jetbrains.python.psi.resolve.PyResolveContext +import com.jetbrains.python.psi.types.TypeEvalContext + + +fun parseDataclassParameters(cls: PyClass, context: TypeEvalContext): DataclassParameters? { + val decorators = cls.decoratorList ?: return null + + for (decorator in decorators.decorators) { + if (PyKnownDecoratorUtil.asKnownDecorators(decorator, context).contains(PyKnownDecoratorUtil.KnownDecorator.DATACLASSES_DATACLASS)) { + for (mapping in decorator.multiMapArguments(PyResolveContext.noImplicits().withTypeEvalContext(context))) { + if (mapping.unmappedArguments.isEmpty() && mapping.unmappedParameters.isEmpty()) { + val builder = DataclassParametersBuilder() + + mapping + .mappedParameters + .entries + .forEach { + builder.update(it.value.name, it.key) + } + + return builder.build() + } + } + } + } + + return null +} + +data class DataclassParameters(val init: Boolean, + val repr: Boolean, + val eq: Boolean, + val order: Boolean, + val hash: Boolean?, + val frozen: Boolean, + val initArgument: PyExpression?, + val reprArgument: PyExpression?, + val eqArgument: PyExpression?, + val orderArgument: PyExpression?, + val hashArgument: PyExpression?, + val frozenArgument: PyExpression?) + +private class DataclassParametersBuilder { + + companion object { + private val DEFAULT_INIT: Boolean = true + private val DEFAULT_REPR: Boolean = true + private val DEFAULT_EQ: Boolean = true + private val DEFAULT_ORDER: Boolean = false + private val DEFAULT_HASH: Boolean? = null // `null` means `None` + private val DEFAULT_FROZEN: Boolean = false + } + + private var init = DEFAULT_INIT + private var repr = DEFAULT_REPR + private var eq = DEFAULT_EQ + private var order = DEFAULT_ORDER + private var hash = DEFAULT_HASH + private var frozen = DEFAULT_FROZEN + + private var initArgument: PyExpression? = null + private var reprArgument: PyExpression? = null + private var eqArgument: PyExpression? = null + private var orderArgument: PyExpression? = null + private var hashArgument: PyExpression? = null + private var frozenArgument: PyExpression? = null + + fun update(name: String?, argument: PyExpression?) { + val value = PyUtil.peelArgument(argument) + + when (name) { + "init" -> { + init = evaluateBoolean(value, DEFAULT_INIT) + initArgument = argument + } + "repr" -> { + repr = evaluateBoolean(value, DEFAULT_REPR) + reprArgument = argument + } + "eq" -> { + eq = evaluateBoolean(value, DEFAULT_EQ) + eqArgument = argument + } + "order" -> { + order = evaluateBoolean(value, DEFAULT_ORDER) + orderArgument = argument + } + "hash" -> { + hash = evaluateBoolean(value, DEFAULT_HASH) + hashArgument = argument + } + "frozen" -> { + frozen = evaluateBoolean(value, DEFAULT_FROZEN) + frozenArgument = argument + } + } + } + + fun build() = DataclassParameters(init, repr, eq, order, hash, frozen, + initArgument, reprArgument, eqArgument, orderArgument, hashArgument, frozenArgument) + + + private inline fun evaluateBoolean(expression: PyExpression?, default: T): T { + val result = PyEvaluator().evaluate(expression) + if (result is Boolean) return result as T + return default + } +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt index 56aecf472540..4fb3ea084eb0 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt @@ -28,7 +28,8 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() { } private fun getDataclassTypeForClass(cls: PyClass, context: TypeEvalContext): PyCallableType? { - if (!PyKnownDecoratorUtil.getKnownDecorators(cls, context).contains(PyKnownDecoratorUtil.KnownDecorator.DATACLASSES_DATACLASS)) { + val dataclassParameters = parseDataclassParameters(cls, context) + if (dataclassParameters == null || !dataclassParameters.init) { return null } diff --git a/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py b/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py index 1097d94e02fb..3132d91ade04 100644 --- a/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py +++ b/python/testData/inspections/PyArgumentListInspection/InitializingDataclass/a.py @@ -10,4 +10,46 @@ A() A(1, "a") A(1, "a", 1.0) -A(1, "a", 1.0, "b") \ No newline at end of file +A(1, "a", 1.0, "b") + + +@dataclasses.dataclass(init=True) +class A2: + x: int + y: str + z: float = 0.0 + +A2() +A2(1) +A2(1, "a") +A2(1, "a", 1.0) +A2(1, "a", 1.0, "b") + + +@dataclasses.dataclass(init=False) +class B1: + x: int = 1 + y: str = "2" + z: float = 0.0 + +B1() +B1(1) +B1(1, "a") +B1(1, "a", 1.0) +B1(1, "a", 1.0, "b") + + +@dataclasses.dataclass(init=False) +class B2: + x: int + y: str + z: float = 0.0 + + def __init__(self, x: int): + self.x = x + self.y = str(x) + self.z = 0.0 + +B2() +B2(1) +B2(1, 2) \ No newline at end of file diff --git a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py index 386645bb2d70..27c8a7480ac7 100644 --- a/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py +++ b/python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py @@ -10,4 +10,48 @@ A(1, "a") A("a", 1) A(1, "a", 1.0) -A("a", 1, "b") \ No newline at end of file +A("a", 1, "b") + + +@dataclasses.dataclass(init=True) +class A2: + x: int + y: str + z: float = 0.0 + +A2(1, "a") +A2("a", 1) + +A2(1, "a", 1.0) +A2("a", 1, "b") + + +@dataclasses.dataclass(init=False) +class B1: + x: int = 1 + y: str = "2" + z: float = 0.0 + +B1(1) +B1("1") + +B1(1, "a") +B1("a", 1) + +B1(1, "a", 1.0) +B1("a", 1, "b") + + +@dataclasses.dataclass(init=False) +class B2: + x: int + y: str + z: float = 0.0 + + def __init__(self, x: int): + self.x = x + self.y = str(x) + self.z = 0.0 + +B2(1) +B2("1") \ No newline at end of file diff --git a/python/testData/paramInfo/InitializingDataclass/a.py b/python/testData/paramInfo/InitializingDataclass/a.py index 61ae24848fae..eb3be3a71599 100644 --- a/python/testData/paramInfo/InitializingDataclass/a.py +++ b/python/testData/paramInfo/InitializingDataclass/a.py @@ -6,4 +6,36 @@ class A: y: str z: float = 0.0 -A() \ No newline at end of file +A() + + +@dataclasses.dataclass(init=True) +class A2: + x: int + y: str + z: float = 0.0 + +A2() + + +@dataclasses.dataclass(init=False) +class B1: + x: int = 1 + y: str = "2" + z: float = 0.0 + +B1() + + +@dataclasses.dataclass(init=False) +class B2: + x: int + y: str + z: float = 0.0 + + def __init__(self, x: int): + self.x = x + self.y = str(x) + self.z = 0.0 + +B2() \ 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 301d7539b9af..fd01d752efc0 100644 --- a/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java +++ b/python/testSrc/com/jetbrains/python/PyParameterInfoTest.java @@ -712,9 +712,18 @@ public class PyParameterInfoTest extends LightMarkedTestCase { runWithLanguageLevel( LanguageLevel.PYTHON37, () -> { - final Map marks = loadMultiFileTest(1); + final Map marks = loadMultiFileTest(4); 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, "}); + + feignCtrlP(marks.get("").getTextOffset()).check( + Arrays.asList("self: object", "cls: object"), + Arrays.asList(ArrayUtil.EMPTY_STRING_ARRAY, ArrayUtil.EMPTY_STRING_ARRAY), + Arrays.asList(new String[]{"self: object"}, new String[]{"cls: object"}) + ); + + feignCtrlP(marks.get("").getTextOffset()).check("self: B2, x: int", new String[]{"x: int"}, new String[]{"self: B2, "}); } ); }