From 6e0ad4e0c8000f84baf0c5fb0b48e6b153bad81d Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Fri, 30 Mar 2018 20:28:07 +0300
Subject: [PATCH] Don't suggest print function if Py2 file does not contain
 corresponding import from future (PY-28103)

---
 .../psi/impl/references/PyReferenceImpl.java  | 19 +++++++++++++------
 .../python/PythonCompletionTest.java          | 14 ++++++++++++++
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
index 4c2fa489536a..a229b7aa3bdc 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
@@ -679,18 +679,25 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     final PyQualifiedExpression element = originalElement != null ? originalElement : myElement;
     final PsiElement realContext = PyPsiUtils.getRealContext(element);
 
-    // include our own names
-    final int underscores = PyUtil.getInitialUnderscores(element.getName());
     final PyBuiltinCache builtinCache = PyBuiltinCache.getInstance(element);
+    final LanguageLevel languageLevel = LanguageLevel.forElement(myElement);
     final CompletionVariantsProcessor processor = new CompletionVariantsProcessor(element, e -> {
       if (builtinCache.isBuiltin(e)) {
-        final String name = e instanceof PyElement ? ((PyElement)e).getName() : null;
         if (e instanceof PyImportElement) {
           return false;
         }
-        if (name != null && PyUtil.getInitialUnderscores(name) == 1) {
+
+        final String name = e instanceof PyElement ? ((PyElement)e).getName() : null;
+        if (PyUtil.getInitialUnderscores(name) == 1) {
           return false;
         }
+
+        if (languageLevel.isPython2() && PyNames.PRINT.equals(name)) {
+          final PyFile file = PyUtil.as(myElement.getContainingFile(), PyFile.class);
+          if (file != null && !file.hasImportFromFuture(FutureFeature.PRINT_FUNCTION)) {
+            return false;
+          }
+        }
       }
       return true;
     }, null);
@@ -709,7 +716,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
       PyResolveUtil.scopeCrawlUp(processor, builtinsFile, null, null);
     }
 
-    if (underscores >= 2) {
+    if (PyUtil.getInitialUnderscores(element.getName()) >= 2) {
       // if we're a normal module, add module's attrs
       if (realContext.getContainingFile() instanceof PyFile) {
         for (String name : PyModuleType.getPossibleInstanceMembers()) {
@@ -718,7 +725,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
       }
 
       // if we're inside method, add implicit __class__
-      if (!LanguageLevel.forElement(myElement).isPython2()) {
+      if (!languageLevel.isPython2()) {
         Optional
           .ofNullable(PsiTreeUtil.getParentOfType(myElement, PyFunction.class))
           .map(PyFunction::getContainingClass)
diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
index e24c215c5c91..9f13f9fd1052 100644
--- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
@@ -1335,6 +1335,20 @@ public class PythonCompletionTest extends PyTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON34, this::assertSingleVariantInExtendedCompletion);
   }
 
+  // PY-28103
+  public void testPrintFunctionWithoutFuture() {
+    final List<String> suggested = doTestByText("pr<caret>");
+    assertNotNull(suggested);
+    assertSameElements(suggested, "print", "property", "repr");
+  }
+
+  // PY-28103
+  public void testPrintFunctionWithFuture() {
+    final List<String> suggested = doTestByText("from __future__ import print_function\npr<caret>");
+    assertNotNull(suggested);
+    assertSameElements(suggested, "print", "print", "print_function", "property", "repr");
+  }
+
   private void assertNoVariantsInExtendedCompletion() {
     myFixture.copyDirectoryToProject(getTestName(true), "");
     myFixture.configureByFile("a.py");