Env tests for tensorflow submodules and subpackages (PY-33034)

GitOrigin-RevId: a98ff02950327a6fa7dc19ec5f88205915281451

python/setup-test-environment/build.gradle

@@ -122,6 +122,18 @@ envs {
                "python3.7",
                true)
 
+  // For TensorFlow
+  createPython("py37_tensorflow1",
+               "3.7.3",
+               ["tensorflow < 2"],
+               "python3.7\npython3\ntensorflow1",
+               true)
+  createPython("py37_tensorflow2",
+               "3.7.3",
+               ["tensorflow ~= 2.0.0b1"],
+               "python3.7\npython3\ntensorflow2",
+               true)
+
   if (Os.isFamily(Os.FAMILY_UNIX)) {
     createPython("qt_env", "3.7.2", ["pyqt5==5.12", "PySide2==5.12.1"], "qt", true)
   }

python/testData/tensorflow/compatv1modules.txt

@@ -0,0 +1,53 @@
+app
+audio
+autograph
+bitwise
+compat
+config
+data
+debugging
+distribute
+distributions
+dtypes
+errors
+experimental
+feature_column
+flags
+gfile
+graph_util
+image
+initializers
+io
+layers
+linalg
+lite
+logging
+lookup
+losses
+manip
+math
+metrics
+nest
+nn
+profiler
+python_io
+quantization
+queue
+ragged
+random
+raw_ops
+resource_loader
+saved_model
+sets
+signal
+sparse
+spectral
+strings
+summary
+sysconfig
+test
+tpu
+train
+user_ops
+version
+xla

python/testData/tensorflow/modules.py

@@ -0,0 +1,8 @@
+import tensorflow
+root = tensorflow.compat.v1  # or tensorflow
+module_type = type(tensorflow)
+print('\n'.join('%s' % (name,)
+                for name in dir(root)
+                if not name.startswith('_')
+                for module in (getattr(root, name),)
+                if type(module) is module_type))

python/testData/tensorflow/v1modules.txt

@@ -0,0 +1,54 @@
+app
+audio
+autograph
+bitwise
+compat
+config
+data
+debugging
+distribute
+distributions
+dtypes
+errors
+estimator
+experimental
+feature_column
+gfile
+graph_util
+image
+initializers
+io
+keras
+layers
+linalg
+lite
+logging
+lookup
+losses
+manip
+math
+metrics
+nest
+nn
+profiler
+python_io
+quantization
+queue
+ragged
+random
+raw_ops
+resource_loader
+saved_model
+sets
+signal
+sparse
+spectral
+strings
+summary
+sysconfig
+test
+tpu
+train
+user_ops
+version
+xla

python/testData/tensorflow/v2modules.txt

@@ -0,0 +1,45 @@
+audio
+autograph
+bitwise
+compat
+config
+data
+debugging
+distribute
+dtypes
+errors
+estimator
+experimental
+feature_column
+graph_util
+image
+initializers
+io
+keras
+linalg
+lite
+lookup
+losses
+math
+metrics
+nest
+nn
+optimizers
+quantization
+queue
+ragged
+random
+raw_ops
+saved_model
+sets
+signal
+sparse
+strings
+summary
+sysconfig
+test
+tools
+tpu
+train
+version
+xla

python/testSrc/com/jetbrains/env/tensorFlow/PyTensorFlowTest.kt

@@ -0,0 +1,83 @@
+// Copyright 2000-2019 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.env.tensorFlow
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.util.Computable
+import com.intellij.psi.PsiElement
+import com.intellij.testFramework.UsefulTestCase
+import com.jetbrains.env.EnvTestTagsRequired
+import com.jetbrains.env.PyEnvTestCase
+import com.jetbrains.env.PyExecutionFixtureTestTask
+import com.jetbrains.python.PythonFileType
+import com.jetbrains.python.PythonTestUtil
+import com.jetbrains.python.psi.PyReferenceExpression
+import com.jetbrains.python.tools.sdkTools.SdkCreationType
+import org.junit.Test
+import java.nio.file.Files
+import java.nio.file.Paths
+
+class PyTensorFlowTest : PyEnvTestCase() {
+
+  @Test
+  @EnvTestTagsRequired(tags = ["tensorflow1"])
+  fun tensorFlow1Modules() {
+    runPythonTest(TensorFlowModulesTask("tensorflow", "v1modules.txt"))
+  }
+
+  @Test
+  @EnvTestTagsRequired(tags = ["tensorflow2"])
+  fun tensorFlow2Modules() {
+    runPythonTest(TensorFlowModulesTask("tensorflow", "v2modules.txt"))
+  }
+
+  @Test
+  @EnvTestTagsRequired(tags = ["tensorflow2"])
+  fun compatTensorFlow1Modules() {
+    runPythonTest(TensorFlowModulesTask("tensorflow.compat.v1", "compatv1modules.txt"))
+  }
+
+  private class TensorFlowModulesTask(private val prefix: String, private val fileName: String) : PyExecutionFixtureTestTask(null) {
+
+    override fun runTestOn(sdkHome: String, existingSdk: Sdk?) {
+      createTempSdk(sdkHome, SdkCreationType.SDK_PACKAGES_AND_SKELETONS)
+
+      val modules = Files.readAllLines(Paths.get(PythonTestUtil.getTestDataPath(), "tensorflow", fileName))
+      runCompletion(modules)
+      runResolve(modules)
+    }
+
+    private fun runCompletion(modules: List<String>) {
+      // `from tensorflow.<m>` completes
+      // `tensorflow.<m>` completes
+      configureAndCompleteAtCaret("from $prefix.<caret>", modules)
+      configureAndCompleteAtCaret("import $prefix\n$prefix.<caret>", modules)
+    }
+
+    private fun runResolve(modules: List<String>) {
+      // `from tensorflow.<m>` resolves
+      // `tensorflow.<m>` resolves
+      // `from tensorflow.<m>` resolves to the same as `tensorflow.<m>`
+
+      modules.forEach {
+        UsefulTestCase.assertSame(
+          configureAndResolveAtCaret("import $prefix.$it<caret>"),
+          configureAndResolveAtCaret("from $prefix.$it<caret> import *")
+        )
+      }
+    }
+
+    private fun configureAndCompleteAtCaret(text: String, modules: List<String>) {
+      myFixture.configureByText(PythonFileType.INSTANCE, text)
+      myFixture.completeBasic()
+      UsefulTestCase.assertContainsElements(myFixture.lookupElementStrings!!, modules)
+    }
+
+    private fun configureAndResolveAtCaret(text: String): PsiElement {
+      myFixture.configureByText(PythonFileType.INSTANCE, text)
+      return ApplicationManager.getApplication().runReadAction(
+        Computable { (myFixture.file.findElementAt(myFixture.caretOffset - 1)!!.parent as PyReferenceExpression).reference.resolve()!! }
+      )
+    }
+  }
+}