diff --git a/python/src/com/jetbrains/python/psi/PyWithItem.java b/python/src/com/jetbrains/python/psi/PyWithItem.java index c9fabb094e2f..2b6998943f0e 100644 --- a/python/src/com/jetbrains/python/psi/PyWithItem.java +++ b/python/src/com/jetbrains/python/psi/PyWithItem.java @@ -1,10 +1,13 @@ package com.jetbrains.python.psi; +import org.jetbrains.annotations.Nullable; + /** * @author yole */ public interface PyWithItem extends PyElement { PyWithItem[] EMPTY_ARRAY = new PyWithItem[0]; - PyTargetExpression getTargetExpression(); + @Nullable + PyExpression getTargetExpression(); } diff --git a/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java b/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java index 492f41b6c5cc..8e1e187e2834 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java @@ -1,9 +1,10 @@ package com.jetbrains.python.psi.impl; import com.intellij.lang.ASTNode; -import com.jetbrains.python.PyElementTypes; -import com.jetbrains.python.psi.PyTargetExpression; +import com.jetbrains.python.PyTokenTypes; +import com.jetbrains.python.psi.PyExpression; import com.jetbrains.python.psi.PyWithItem; +import org.jetbrains.annotations.Nullable; /** * @author yole @@ -13,9 +14,18 @@ public class PyWithItemImpl extends PyElementImpl implements PyWithItem { super(astNode); } - public PyTargetExpression getTargetExpression() { - final ASTNode asNameNode = getNode().findChildByType(PyElementTypes.TARGET_EXPRESSION); - if (asNameNode == null) return null; - return (PyTargetExpression)asNameNode.getPsi(); + @Nullable + public PyExpression getTargetExpression() { + ASTNode[] children = getNode().getChildren(null); + boolean foundAs = false; + for (ASTNode child : children) { + if (child.getElementType() == PyTokenTypes.AS_KEYWORD) { + foundAs = true; + } + else if (foundAs && child.getPsi() instanceof PyExpression) { + return (PyExpression) child.getPsi(); + } + } + return null; } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java index e5804887d567..7e3ace0beac2 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java @@ -5,10 +5,7 @@ import com.intellij.psi.PsiElement; import com.intellij.psi.tree.TokenSet; import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.PyElementTypes; -import com.jetbrains.python.psi.PyElement; -import com.jetbrains.python.psi.PyElementVisitor; -import com.jetbrains.python.psi.PyWithItem; -import com.jetbrains.python.psi.PyWithStatement; +import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; import java.util.ArrayList; @@ -34,7 +31,8 @@ public class PyWithStatementImpl extends PyElementImpl implements PyWithStatemen List result = new ArrayList(); if (items != null) { for (PyWithItem item : items) { - result.add(item.getTargetExpression()); + PyExpression targetExpression = item.getTargetExpression(); + result.addAll(PyUtil.flattenedParens(targetExpression)); } } return result; diff --git a/python/testData/resolve/WithParentheses.py b/python/testData/resolve/WithParentheses.py new file mode 100644 index 000000000000..8e9ec981c9ae --- /dev/null +++ b/python/testData/resolve/WithParentheses.py @@ -0,0 +1,6 @@ +from contextlib import nested +with nested(patch('Package.ModuleName.ClassName'), + patch('Package.ModuleName.ClassName2', TestUtils.MockClass2)) as (MockClass1, MockClass2): + MockClass1.test.return_value = True +# + MockClass2.test.return_value = True \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java index 5154e86d6a09..5508a00578a1 100644 --- a/python/testSrc/com/jetbrains/python/PyResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyResolveTest.java @@ -384,4 +384,8 @@ public class PyResolveTest extends PyResolveTestCase { final PsiElement element = doResolve(); assertNull(element); } + + public void testWithParentheses() { + assertResolvesTo(LanguageLevel.PYTHON27, PyTargetExpression.class, "MockClass1"); + } } \ No newline at end of file