PY-21479 Handle references in f-strings used in parts of comprehension other than its result

This commit is contained in:
Mikhail Golubev
2016-12-28 15:56:23 +03:00
parent 650ea52539
commit e983d60050
5 changed files with 23 additions and 11 deletions

View File

@@ -41,8 +41,6 @@ import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.List;
import static com.jetbrains.python.psi.PyUtil.as;
/**
* User : ktisha
*/
@@ -101,19 +99,19 @@ public class PyDocReference extends PyReferenceImpl {
@Nullable
private PsiElement getScopeControlFlowAnchor(@NotNull PsiLanguageInjectionHost host) {
if (isInsideFormattedStringNode(host)) {
// Comprehension's result expression is preserved in CFG: this anchor is necessary for flow-sensitive getResultsFromProcessor()
final PsiElement parentComprehensionResult = PsiTreeUtil.findFirstParent(host, PyDocReference::isComprehensionResult);
if (parentComprehensionResult != null) {
return parentComprehensionResult;
final PsiElement comprehensionPart = PsiTreeUtil.findFirstParent(host, PyDocReference::isComprehensionResultOrComponent);
if (comprehensionPart != null) {
return comprehensionPart;
}
return PsiTreeUtil.getParentOfType(host, PyStatement.class);
}
return null;
}
private static boolean isComprehensionResult(@NotNull PsiElement element) {
final PyComprehensionElement comprehension = as(element.getParent(), PyComprehensionElement.class);
return comprehension != null && comprehension.getResultExpression() == element;
private static boolean isComprehensionResultOrComponent(@NotNull PsiElement element) {
// Any comprehension component and its result are represented as children expressions of the comprehension element.
// Only they have respective nodes in CFG and thus can be used as anchors for getResultsFromProcessor().
return element instanceof PyExpression && element.getParent() instanceof PyComprehensionElement;
}
private boolean isInsideFormattedStringNode(@NotNull PsiLanguageInjectionHost host) {

View File

@@ -0,0 +1,2 @@
[42 for foo in range(10) if f'{foo}']
<ref>

View File

@@ -1,2 +1,2 @@
xs = [[x for x in f'{foo}'] for foo in range(10)]
<ref>
xs = [42 for foo in range(10) for _ in f'{foo}']
<ref>

View File

@@ -0,0 +1,2 @@
xs = [[x for x in f'{foo}'] for foo in range(10)]
<ref>

View File

@@ -75,6 +75,16 @@ public class PyInjectionResolveTest extends PyResolveTestCase {
assertResolvesTo(LanguageLevel.PYTHON36, PyTargetExpression.class, "foo");
}
// PY-21479
public void testFStringNestedInResultComprehensionSourcePart() {
assertResolvesTo(LanguageLevel.PYTHON36, PyTargetExpression.class, "foo");
}
// PY-21479
public void testFStringComprehensionConditionPart() {
assertResolvesTo(LanguageLevel.PYTHON36, PyTargetExpression.class, "foo");
}
// PY-21479
public void testFStringNestedComprehensionSourcePart() {
assertResolvesTo(LanguageLevel.PYTHON36, PyTargetExpression.class, "foo");