diff --git a/python/src/com/jetbrains/python/inspections/PyTupleAssignmentBalanceInspection.java b/python/src/com/jetbrains/python/inspections/PyTupleAssignmentBalanceInspection.java
index 694b366ff9a9..d773ac9a91fc 100644
--- a/python/src/com/jetbrains/python/inspections/PyTupleAssignmentBalanceInspection.java
+++ b/python/src/com/jetbrains/python/inspections/PyTupleAssignmentBalanceInspection.java
@@ -3,10 +3,10 @@ package com.jetbrains.python.inspections;
import com.intellij.codeInspection.LocalInspectionToolSession;
import com.intellij.codeInspection.ProblemsHolder;
import com.intellij.openapi.vfs.VirtualFile;
-import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiElementVisitor;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyType;
import org.jetbrains.annotations.Nls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
@@ -38,21 +38,15 @@ public class PyTupleAssignmentBalanceInspection extends PyInspection {
@Override
public void visitPyAssignmentStatement(PyAssignmentStatement node) {
PyExpression lhsExpression = node.getLeftHandSideExpression();
- PsiElement assignedValue = node.getAssignedValue();
- if (assignedValue instanceof PyParenthesizedExpression) // PY-2659
- assignedValue = ((PyParenthesizedExpression)assignedValue).getContainedExpression();
+ PyExpression assignedValue = node.getAssignedValue();
if (lhsExpression instanceof PyParenthesizedExpression) // PY-4360
lhsExpression = ((PyParenthesizedExpression)lhsExpression).getContainedExpression();
- if (assignedValue instanceof PyReferenceExpression) { // PY-4357
- assignedValue = ((PyReferenceExpression)assignedValue).followAssignmentsChain(resolveWithoutImplicits()).getElement();
- }
- if (lhsExpression instanceof PyTupleExpression) {
- int valuesLength = 1;
+
+ if (assignedValue == null) return;
+ PyType type = assignedValue.getType(myTypeEvalContext);
+ if (lhsExpression instanceof PyTupleExpression && type != null) {
+ int valuesLength = PyUtil.getElementsCount(assignedValue, myTypeEvalContext);
PyExpression[] elements = ((PyTupleExpression) lhsExpression).getElements();
- if (assignedValue instanceof PySequenceExpression)
- valuesLength = ((PySequenceExpression)assignedValue).getElements().length;
- else if (assignedValue instanceof PyDictLiteralExpression)
- valuesLength = ((PyDictLiteralExpression)assignedValue).getElements().length;
boolean containsStarExpression = false;
VirtualFile virtualFile = node.getContainingFile().getVirtualFile();
@@ -71,9 +65,9 @@ public class PyTupleAssignmentBalanceInspection extends PyInspection {
int targetsLength = elements.length;
if (targetsLength > valuesLength) {
- registerProblem(node.getAssignedValue(), "Need more values to unpack");
+ registerProblem(assignedValue, "Need more values to unpack");
} else if (!containsStarExpression && targetsLength < valuesLength) {
- registerProblem(node.getAssignedValue(), "Too many values to unpack");
+ registerProblem(assignedValue, "Too many values to unpack");
}
}
}
diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java
index 7595b935b5b2..c4f962d39c11 100644
--- a/python/src/com/jetbrains/python/psi/PyUtil.java
+++ b/python/src/com/jetbrains/python/psi/PyUtil.java
@@ -38,6 +38,7 @@ import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.documentation.EpydocUtil;
import com.jetbrains.python.psi.impl.PyBuiltinCache;
import com.jetbrains.python.psi.types.PyClassType;
+import com.jetbrains.python.psi.types.PyTupleType;
import com.jetbrains.python.psi.types.PyType;
import com.jetbrains.python.psi.types.TypeEvalContext;
import com.jetbrains.python.refactoring.classes.extractSuperclass.PyExtractSuperclassHelper;
@@ -1044,5 +1045,24 @@ public class PyUtil {
}
return (PyFile)psi;
}
+
+ /**
+ * counts elements in iterable
+ * @param expression to count containing elements (iterable)
+ * @return element count
+ */
+ public static int getElementsCount(PyExpression expression, TypeEvalContext evalContext) {
+ int valuesLength = 1;
+ PyType type = expression.getType(evalContext);
+ if (type instanceof PyTupleType)
+ valuesLength = ((PyTupleType)type).getElementCount();
+ if (expression instanceof PySequenceExpression)
+ valuesLength = ((PySequenceExpression)expression).getElements().length;
+ else if (expression instanceof PyDictLiteralExpression)
+ valuesLength = ((PyDictLiteralExpression)expression).getElements().length;
+ else if (expression instanceof PyStringLiteralExpression)
+ valuesLength = ((PyStringLiteralExpression)expression).getStringValue().length();
+ return valuesLength;
+ }
}
diff --git a/python/testData/inspections/PyTupleAssignmentBalanceInspection3/test.py b/python/testData/inspections/PyTupleAssignmentBalanceInspection3/test.py
index 7113ee884c6f..64a709385a1e 100644
--- a/python/testData/inspections/PyTupleAssignmentBalanceInspection3/test.py
+++ b/python/testData/inspections/PyTupleAssignmentBalanceInspection3/test.py
@@ -10,6 +10,6 @@ a, b = c
#PY-4358
a, b = [1, 2, 3]
-a, b = 'str'
+a, b = 'str'
a, b = {1, 2, 3}
a, b = {1:2, 2: 3, 3:4}
\ No newline at end of file