PY-38581 Implement previous calls feature

GitOrigin-RevId: 5682350bfd95a6abfc325b511c55d6e27b41476a
2026-09-10 13:17:09 +07:00 · 2020-02-25 05:37:15 +00:00
parent b2d1270096
commit 073af04a6a
30 changed files with 651 additions and 88 deletions
@@ -616,6 +616,7 @@ class CommunityLibraryLicenses {
    jetbrainsLibrary("kotlin-stdlib-jdk8"),
    jetbrainsLibrary("kotlin-test"),
    jetbrainsLibrary("kotlinx-coroutines-jdk8"),
+    jetbrainsLibrary("ml-completion-prev-exprs-models"),
    jetbrainsLibrary("precompiled_jshell-frontend"),
    jetbrainsLibrary("rd-core"),
    jetbrainsLibrary("rd-framework"),
@@ -51,5 +51,21 @@
    </orderEntry>
    <orderEntry type="library" name="StreamEx" level="project" />
    <orderEntry type="module" module-name="intellij.platform.statistics" />
+    <orderEntry type="module-library">
+      <library name="ml-completion-prev-exprs-models" type="repository">
+        <properties maven-id="completion.ml.python.features:ml-completion-prev-exprs-models:1.11" />
+        <CLASSES>
+          <root url="jar://$MAVEN_REPOSITORY$/completion/ml/python/features/ml-completion-prev-exprs-models/1.11/ml-completion-prev-exprs-models-1.11.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/kotlin/kotlin-stdlib-jdk8/1.3.41/kotlin-stdlib-jdk8-1.3.41.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/kotlin/kotlin-stdlib/1.3.41/kotlin-stdlib-1.3.41.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/kotlin/kotlin-stdlib-common/1.3.41/kotlin-stdlib-common-1.3.41.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/annotations/13.0/annotations-13.0.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/kotlin/kotlin-stdlib-jdk7/1.3.41/kotlin-stdlib-jdk7-1.3.41.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/com/google/code/gson/gson/2.8.5/gson-2.8.5.jar!/" />
+        </CLASSES>
+        <JAVADOC />
+        <SOURCES />
+      </library>
+    </orderEntry>
  </component>
 </module>
@@ -4,6 +4,8 @@ package com.jetbrains.python.codeInsight.mlcompletion
 import com.intellij.codeInsight.completion.ml.CompletionEnvironment
 import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PrevCallsModelsStorage
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PyPrevCallsCompletionFeatures
 import com.jetbrains.python.psi.types.TypeEvalContext

 class PyContextFeatureProvider : ContextFeatureProvider {
@@ -58,6 +60,14 @@ class PyContextFeatureProvider : ContextFeatureProvider {
      result["containing_class_have_constructor"] = MLFeatureValue.binary(classHaveConstructor)
    }}

+    val cursorOffset = environment.lookup.lookupStart
+    val isInCondition = result["is_in_condition"]?.value as? Boolean ?: false
+    val isInForStatement = result["is_in_for_statement"]?.value as? Boolean ?: false
+    PyPrevCallsCompletionFeatures.calculatePrevCallsContextInfo(cursorOffset, position, isInCondition, isInForStatement)?.let {
+      PrevCallsModelsStorage.loadModelFor(it.qualifier)
+      environment.putUserData(PyPrevCallsCompletionFeatures.PREV_CALLS_CONTEXT_INFO_KEY, it)
+    }
+
    return result
  }
 }
@@ -6,6 +6,7 @@ import com.intellij.codeInsight.completion.ml.ContextFeatures
 import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PyPrevCallsCompletionFeatures

 class PyElementFeatureProvider : ElementFeatureProvider {
  override fun getName(): String = "python"
@@ -77,6 +78,17 @@ class PyElementFeatureProvider : ElementFeatureProvider {
      result["receiver_tokens_num"] = MLFeatureValue.numerical(receiverTokensNum)
    }}

+    contextFeatures.getUserData(PyPrevCallsCompletionFeatures.PREV_CALLS_CONTEXT_INFO_KEY)?.let { contextInfo ->
+      PyPrevCallsCompletionFeatures.getResult(lookupString, contextInfo)?.let { with (it) {
+        when {
+          primaryWeight != null -> result["prev_2_calls_primary_weight"] = MLFeatureValue.numerical(primaryWeight!!)
+          weightOneCall != null -> result["prev_2_calls_weight_one_call"] = MLFeatureValue.numerical(weightOneCall!!)
+          weightSecondaryTwoCalls != null -> result["prev_2_calls_weight_secondary_two_calls"] = MLFeatureValue.numerical(weightSecondaryTwoCalls!!)
+          weightEmptyPrevCalls != null -> result["prev_2_calls_weight_empty_prev_calls"] = MLFeatureValue.numerical(weightEmptyPrevCalls!!)
+        }
+      }}
+    }
+
    return result
  }
 }
@@ -0,0 +1,85 @@
+// 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.python.codeInsight.mlcompletion.prev2calls
+
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.psi.*
+
+class AssignmentVisitor(private val borderOffset: Int,
+                        private val scope: PsiElement,
+                        private val fullNames: MutableMap<String, String>) : PyRecursiveElementVisitor() {
+  data class QualifierAndReference(val qualifier: String, val reference: String)
+
+  val arrPrevCalls = ArrayList<QualifierAndReference>()
+
+  override fun visitPyReferenceExpression(node: PyReferenceExpression) {
+    super.visitPyReferenceExpression(node)
+    if (node.textOffset > borderOffset) return
+
+    val (resolvedExpression, resolvedPrefix) = getResolvedExpression(node)
+    if (node.parent !is PyCallExpression && ("." !in resolvedExpression || resolvedPrefix == resolvedExpression)) return
+
+    val qualifier = resolvedExpression.substringBeforeLast(".", "")
+    val reference = resolvedExpression.substringAfterLast(".")
+    arrPrevCalls.add(QualifierAndReference(qualifier, reference))
+  }
+
+  override fun visitPyWithStatement(node: PyWithStatement) {
+    if (node.textOffset > borderOffset) return
+    node.withItems.filter { it.expression != null && it.target != null }.forEach {
+      fullNames[it.target!!.text] = getResolvedExpression(it.expression).resolvedExpression
+    }
+    super.visitPyWithStatement(node)
+  }
+
+  override fun visitPyAssignmentStatement(node: PyAssignmentStatement) {
+    if (node.textOffset > borderOffset) return
+    super.visitPyAssignmentStatement(node)
+
+    node.targetsToValuesMapping.forEach {
+      val left = it.first
+      val right = it.second
+      if (left is PyTargetExpression) {
+        val leftName = getQualifierComponents(left).joinToString(".")
+        val rightName = getResolvedExpression(right).resolvedExpression
+        if (rightName.isNotEmpty() && leftName.isNotEmpty()) {
+          fullNames[leftName] = rightName
+        }
+      }
+    }
+  }
+
+  data class ResolvedExpression(val resolvedExpression: String = "", val resolvedPrefix: String = "")
+  private fun getResolvedExpression(node: PsiElement?): ResolvedExpression {
+    if (node == null) return ResolvedExpression()
+
+    val components = getQualifierComponents(node)
+    for (i in components.indices) {
+      val firstN = i + 1
+      val prefix = components.take(firstN).joinToString(".")
+      fullNames[prefix]?.let { resolvedPrefix ->
+        val postfix =
+          if (firstN < components.size)
+            components.takeLast(components.size - firstN).joinToString(separator=".", prefix=".")
+          else ""
+        return ResolvedExpression("$resolvedPrefix$postfix", resolvedPrefix)
+      }
+    }
+
+    return ResolvedExpression(components.joinToString("."))
+  }
+
+  private fun getQualifierComponents(node: PsiElement): List<String> =
+    generateSequence(node) { it.firstChild }
+      .filterIsInstance<PyQualifiedExpression>()
+      .mapNotNull { it.name }
+      .toList()
+      .asReversed()
+
+  override fun visitPyFunction(node: PyFunction) {
+    if (node == scope) super.visitPyFunction(node)
+  }
+
+  override fun visitPyClass(node: PyClass) {
+    if (node == scope) super.visitPyClass(node)
+  }
+}
@@ -0,0 +1,31 @@
+// 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.python.codeInsight.mlcompletion.prev2calls
+
+import com.jetbrains.python.psi.PyFromImportStatement
+import com.jetbrains.python.psi.PyImportStatement
+import com.jetbrains.python.psi.PyRecursiveElementVisitor
+
+class ImportsVisitor(val fullNames: MutableMap<String, String> = mutableMapOf()): PyRecursiveElementVisitor() {
+  override fun visitPyFromImportStatement(node: PyFromImportStatement?) {
+    super.visitPyFromImportStatement(node)
+    if (node == null) return
+
+    val fromName = node.importSourceQName
+    node.importElements.forEach { importElement ->
+      val importedQName = importElement.importedQName.toString()
+      val fullName = "$fromName.$importedQName"
+      fullNames[importedQName] = fullName
+      importElement.asName?.let { fullNames[it] = fullName }
+    }
+  }
+
+  override fun visitPyImportStatement(node: PyImportStatement?) {
+    super.visitPyImportStatement(node)
+    if (node == null) return
+
+    node.importElements.forEach { importElement ->
+      val importedQName = importElement.importedQName.toString()
+      importElement.asName?.let { fullNames[it] = importedQName }
+    }
+  }
+}
@@ -0,0 +1,65 @@
+// 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.python.codeInsight.mlcompletion.prev2calls
+
+import com.completion.features.models.prevcalls.python.PrevCallsModel
+import com.completion.features.models.prevcalls.python.PrevCallsModelsLoader
+import com.google.common.cache.CacheBuilder
+import com.google.common.cache.CacheLoader
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.util.ui.update.MergingUpdateQueue
+import com.intellij.util.ui.update.Update
+import java.util.concurrent.ExecutionException
+
+object PrevCallsModelsStorage {
+  internal class ModelNotFoundException: Exception()
+  private val logger = Logger.getInstance(PrevCallsModelsStorage::class.java)
+
+  private val package2model = CacheBuilder.newBuilder()
+    .softValues()
+    .maximumSize(40)
+    .build(object: CacheLoader<String, PrevCallsModel>() {
+      override fun load(expression: String): PrevCallsModel? {
+        val result = PrevCallsModelsLoader.getModelForExpression(expression)
+        if (result == null) throw ModelNotFoundException()
+        return result
+      }
+    })
+
+  private val modelsLoadingQueue = MergingUpdateQueue("ModelsLoadingQueue", 1000, true, null,
+                                                      null, null, false)
+  private fun createUpdate(identity: Any, runnable: () -> Unit) = object : Update(identity) {
+    override fun canEat(update: Update?) = this == update
+    override fun run() = runnable()
+  }
+
+  fun loadModelFor(qualifierName: String) {
+    val moduleName = qualifierName.substringBefore(".")
+    if (!haveModelForQualifier(moduleName)) return
+
+    val modelForPackage = package2model.getIfPresent(moduleName)
+
+    fun tryLoadModel() {
+      try {
+        package2model.get(moduleName)
+      } catch (ex: ExecutionException) {
+        if (ex.cause !is ModelNotFoundException) {
+          logger.error(ex)
+        }
+      }
+    }
+
+    if (modelForPackage == null) {
+      if (ApplicationManager.getApplication().isUnitTestMode) {
+        tryLoadModel()
+      }
+      else {
+        modelsLoadingQueue.queue(createUpdate(moduleName) { tryLoadModel() })
+      }
+    }
+  }
+
+  fun getModelFor(qualifierName: String) = package2model.getIfPresent(qualifierName)
+
+  fun haveModelForQualifier(qualifier: String) = PrevCallsModelsLoader.haveModule(qualifier.substringBefore("."))
+}
@@ -0,0 +1,45 @@
+// 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.python.codeInsight.mlcompletion.prev2calls
+
+import com.completion.features.models.prevcalls.python.PrevCallsContextInfo
+import com.completion.features.models.prevcalls.python.PrevCallsModelResponse
+import com.intellij.openapi.util.Key
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.PyFunction
+import kotlin.streams.toList
+
+object PyPrevCallsCompletionFeatures {
+  val PREV_CALLS_CONTEXT_INFO_KEY = Key<PrevCallsContextInfo>("py.ml.completion.prev.calls.user.data")
+
+  fun calculatePrevCallsContextInfo(cursorOffset: Int, psiPosition: PsiElement, isInIf: Boolean, isInLoop: Boolean): PrevCallsContextInfo? {
+    val scopePsiElement =
+      PsiTreeUtil.getParentOfType(psiPosition, PyFunction::class.java, PyClass::class.java, PyFile::class.java) ?: return null
+
+    val importsVisitor = ImportsVisitor()
+    psiPosition.containingFile.accept(importsVisitor)
+    val assignmentVisitor = AssignmentVisitor(cursorOffset, scopePsiElement, importsVisitor.fullNames)
+
+    scopePsiElement.accept(assignmentVisitor)
+
+    val allCalls = assignmentVisitor.arrPrevCalls.asReversed()
+    if (allCalls.isEmpty()) return null
+
+    val qualifierName = allCalls[0].qualifier
+    val previousCalls = allCalls
+      .asSequence()
+      .drop(1)
+      .filter { it.qualifier == qualifierName }
+      .map { it.reference }
+      .toList()
+
+    return PrevCallsContextInfo(previousCalls, qualifierName, isInIf, isInLoop)
+  }
+
+  fun getResult(lookupString: String, contextInfo: PrevCallsContextInfo): PrevCallsModelResponse? {
+    val model = PrevCallsModelsStorage.getModelFor(contextInfo.qualifier.substringBefore(".")) ?: return null
+    return model.getWeightForElement(lookupString, contextInfo)
+  }
+}
@@ -0,0 +1,3 @@
+def foo(elem):
+  elem.call1(elem.field1).ref1.call2(elem.field2, 2).ref2.ref3.call3(2)
+  e<caret>
@@ -0,0 +1,4 @@
+import numpy as np
+a = np.arange(np.maximum(1, 2)).reshape(np.minimum(1, 2), np.absolute(-3))
+a.reshape(3, 4)
+<caret>
@@ -0,0 +1,9 @@
+import pandas as pd
+
+def foo:
+  pass
+
+df = pd.read_csv(pd.compat.StringIO(fruit_price))
+df['label'] = df.apply(foo, axis=1)
+
+<caret>
@@ -0,0 +1,12 @@
+import tensorflow as tf
+import numpy as np
+
+labels_map = [np.argmax(c) for c in counts]
+labels_map = tf.convert_to_tensor(labels_map)
+
+cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx)
+
+correct_prediction = tf.equal(cluster_label, tf.cast(tf.argmax(Y, 1), tf.int32))
+accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
+
+<caret>
@@ -0,0 +1,6 @@
+import numpy as np
+class Clzz:
+  def __init__(self):
+    self.a = np.arange(3)
+    self.a.reshape(3, 3)
+    s<caret>
@@ -0,0 +1,3 @@
+import numpy as np
+a = np.arange(42)
+<caret>
@@ -0,0 +1,9 @@
+import tensorflow as tf
+from tensorflow.contrib.factorization import KMeans
+
+kmeans = KMeans(inputs=tf.placeholder(tf.float32, shape=[None, 42]), num_clusters=42, distance_metric='cosine', use_mini_batch=True)
+
+training_graph = kmeans.training_graph()
+print(training_graph)
+
+<caret>
@@ -0,0 +1,11 @@
+import numpy as np
+
+from tensorflow.examples.tutorials.mnist import input_data
+mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
+counts = np.zeros(shape=(42, 42))
+for i in range(42):
+    counts[idx[i]] += mnist.train.labels[i]
+
+<caret>
+
+
@@ -0,0 +1,7 @@
+import numpy as np
+import matplotlib.pyplot as plt
+mu, sigma = 2, 0.5
+v = np.random.normal(mu, sigma, 10000)
+plt.hist(v, bins=50, density=1)
+plt.show()
+<caret>
@@ -0,0 +1,4 @@
+import numpy as np
+a = np.arange(42)
+b = a.reshape(6, 7)
+<caret>
@@ -0,0 +1,2 @@
+import os
+os.environ.<caret>
@@ -0,0 +1,3 @@
+import os
+v = os.environ.get('SOME_KEY')
+if v.<caret>
@@ -0,0 +1,3 @@
+import os
+os.sdfsdfsdf1231
+os.makedirs(os.<caret>)
@@ -0,0 +1,2 @@
+import os
+os.makedirs(os.<caret>)
@@ -0,0 +1,3 @@
+import os
+os.dfsdfsdfsd234234
+os.<caret>
@@ -0,0 +1,2 @@
+import os
+os.<caret>
@@ -0,0 +1,6 @@
+import sys
+
+lns = sys.stdin.readline()
+lns.rstrip()
+lns.rstrip()
+lns.<caret>
@@ -0,0 +1,4 @@
+import sys
+
+for line in sys.stdin.<caret>:
+  print(line)
@@ -0,0 +1,115 @@
+// 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.python.codeInsight.mlcompletion
+
+import com.intellij.codeInsight.lookup.impl.LookupImpl
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.AssignmentVisitor
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.ImportsVisitor
+import com.jetbrains.python.fixtures.PyTestCase
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.PyFunction
+import junit.framework.TestCase
+
+class AssignmentVisitorTest: PyTestCase() {
+  override fun getTestDataPath(): String = super.getTestDataPath() + "/codeInsight/mlcompletion/prev2calls"
+
+  override fun setUp() {
+    super.setUp()
+    setLanguageLevel(LanguageLevel.PYTHON36)
+  }
+
+  fun testAssignmentVisitorOnePrevCall() = testAssignmentVisitor(arrayListOf(Pair("numpy", "arange")))
+  fun testAssignmentVisitorTwoPrevCalls() = testAssignmentVisitor(arrayListOf(
+    Pair("numpy", "arange"),
+    Pair("numpy.arange", "reshape")))
+
+  fun testAssignmentVisitorTwoDifferentPackages() = testAssignmentVisitor(arrayListOf(
+    Pair("matplotlib", "pyplot"),
+    Pair("numpy", "random"),
+    Pair("numpy.random", "normal"),
+    Pair("matplotlib.pyplot", "hist"),
+    Pair("matplotlib.pyplot", "show")))
+
+  fun testAssignmentVisitorTwoDifferentImportTypes() = testAssignmentVisitor(arrayListOf(
+    Pair("tensorflow", "contrib"),
+    Pair("tensorflow.contrib", "factorization"),
+    Pair("tensorflow.contrib.factorization", "KMeans"),
+    Pair("tensorflow", "placeholder"),
+    Pair("tensorflow", "float32"),
+    Pair("tensorflow.contrib.factorization.KMeans", "training_graph"),
+    Pair("", "print")))
+
+  fun testAssignmentVisitorTwoDifferentKindOfImportsAndPackages() = testAssignmentVisitor(arrayListOf(
+    Pair("tensorflow", "examples"),
+    Pair("tensorflow.examples", "tutorials"),
+    Pair("tensorflow.examples.tutorials", "mnist"),
+    Pair("tensorflow.examples.tutorials.mnist.input_data", "read_data_sets"),
+    Pair("numpy", "zeros"),
+    Pair("", "range"),
+    Pair("tensorflow.examples.tutorials.mnist.input_data.read_data_sets", "train"),
+    Pair("tensorflow.examples.tutorials.mnist.input_data.read_data_sets.train", "labels")))
+
+  fun testAssignmentVisitorCheckInArguments() = testAssignmentVisitor(arrayListOf(
+    Pair("numpy", "argmax"),
+    Pair("tensorflow", "convert_to_tensor"),
+    Pair("tensorflow", "nn"),
+    Pair("tensorflow.nn", "embedding_lookup"),
+    Pair("tensorflow", "equal"),
+    Pair("tensorflow", "cast"),
+    Pair("tensorflow", "argmax"),
+    Pair("tensorflow", "int32"),
+    Pair("tensorflow", "reduce_mean"),
+    Pair("tensorflow", "cast"),
+    Pair("tensorflow", "float32")))
+
+  fun testAssignmentVisitorCheckAnotherPackage() = testAssignmentVisitor(arrayListOf(
+    Pair("pandas", "read_csv"),
+    Pair("pandas", "compat"),
+    Pair("pandas.compat", "StringIO"),
+    Pair("pandas.read_csv", "apply")))
+
+  fun testAssignmentVisitorCallArgumentsOrder() = testAssignmentVisitor(arrayListOf(
+    Pair("numpy", "arange"),
+    Pair("numpy", "maximum"),
+    Pair("numpy.arange", "reshape"),
+    Pair("numpy", "minimum"),
+    Pair("numpy", "absolute"),
+    Pair("numpy.arange.reshape", "reshape")))
+
+  fun testAssignmentVisitorClassSelfFields() = testAssignmentVisitor(arrayListOf(
+    Pair("numpy", "arange"),
+    Pair("numpy.arange", "reshape")))
+
+  fun testAssignmentVisitorCallAndReferenceArgumentsOrder() = testAssignmentVisitor(arrayListOf(
+    Pair("elem", "call1"),
+    Pair("elem", "field1"),
+    Pair("elem.call1", "ref1"),
+    Pair("elem.call1.ref1", "call2"),
+    Pair("elem", "field2"),
+    Pair("elem.call1.ref1.call2", "ref2"),
+    Pair("elem.call1.ref1.call2.ref2", "ref3"),
+    Pair("elem.call1.ref1.call2.ref2.ref3", "call3")))
+
+  private fun testAssignmentVisitor(expectedPrevCalls: ArrayList<Pair<String, String>>) {
+    val lookup = invokeCompletionAndGetLookup()
+    val scope =
+      PsiTreeUtil.getParentOfType(lookup.psiElement!!, PyFunction::class.java, PyClass::class.java, PyFile::class.java)!!
+    val importsVisitor = ImportsVisitor()
+    lookup.psiFile!!.accept(importsVisitor)
+    val assignmentsVisitor = AssignmentVisitor(lookup.lookupStart, scope, importsVisitor.fullNames)
+    scope.accept(assignmentsVisitor)
+    val actualPrevCalls = ArrayList<Pair<String, String>>()
+    assignmentsVisitor.arrPrevCalls.forEachIndexed { i, it ->
+      actualPrevCalls.add(Pair(it.qualifier, it.reference))
+    }
+    TestCase.assertEquals(expectedPrevCalls, actualPrevCalls)
+  }
+
+  private fun invokeCompletionAndGetLookup(): LookupImpl {
+    myFixture.configureByFile(getTestName(true) + ".py")
+    myFixture.completeBasic()
+    return myFixture.lookup as LookupImpl
+  }
+}
@@ -1,16 +1,10 @@
 // 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.python.codeInsight.mlcompletion

-import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
-import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
-import com.intellij.codeInsight.lookup.LookupElement
-import com.jetbrains.python.PythonLanguage
-import com.jetbrains.python.fixtures.PyTestCase
 import com.jetbrains.python.psi.LanguageLevel
-import org.junit.Assert

-class PyMlCompletionFeaturesTest: PyTestCase() {
+class PyMlCompletionFeaturesTest: PyMlCompletionTestCase() {
  override fun getTestDataPath(): String = super.getTestDataPath() + "/codeInsight/mlcompletion"

  override fun setUp() {
@@ -251,85 +245,4 @@ class PyMlCompletionFeaturesTest: PyTestCase() {
                                                          Pair("receiver_name_matches", MLFeatureValue.binary(false)),
                                                          Pair("receiver_num_matched_tokens", MLFeatureValue.numerical(3)),
                                                          Pair("receiver_tokens_num", MLFeatureValue.numerical(3)))
-
-  private fun doContextFeaturesTest(vararg expected: Pair<String, MLFeatureValue>) = doContextFeaturesTest(listOf(*expected), emptyList())
-
-  private fun doContextFeaturesTest(expectedDefined: List<Pair<String, MLFeatureValue>>, expectedUndefined: List<String>) {
-    doWithInstalledProviders { contextFeaturesProvider, _ ->
-      invokeCompletion()
-      assertHasFeatures(contextFeaturesProvider.features, expectedDefined)
-      assertHasNotFeatures(contextFeaturesProvider.features, expectedUndefined)
-    }
-  }
-
-  private fun doElementFeaturesTest(checks: List<Pair<String, List<Pair<String, MLFeatureValue>>>>) {
-    checks.forEach {
-      doElementFeaturesTest(it.first, it.second, emptyList())
-    }
-  }
-
-  private fun doElementFeaturesTest(elementToSelect: String, vararg expected: Pair<String, MLFeatureValue>) {
-    doElementFeaturesTest(elementToSelect, arrayListOf(*expected), emptyList())
-  }
-
-  private fun doElementFeaturesTest(elementToSelect: String,
-                                    expectedDefined: List<Pair<String, MLFeatureValue>>,
-                                    expectedUndefined: List<String>) {
-    val selector: (LookupElement) -> Boolean = { it.lookupString == elementToSelect }
-    doElementFeaturesInternalTest(selector, expectedDefined, expectedUndefined)
-  }
-
-  private fun doElementFeaturesInternalTest(selector: (LookupElement) -> Boolean,
-                                            expectedDefined: List<Pair<String, MLFeatureValue>>,
-                                            expectedUndefined: List<String>) {
-    doWithInstalledProviders { _, elementFeaturesProvider ->
-      invokeCompletion()
-
-      val selected = myFixture.lookupElements!!.find(selector)
-      assertNotNull(selected)
-
-      val features = elementFeaturesProvider.features[selected]
-      assertNotNull(features)
-      assertHasFeatures(features!!, expectedDefined)
-      assertHasNotFeatures(features, expectedUndefined)
-    }
-  }
-
-  private fun doWithInstalledProviders(action: (contextFeaturesProvider: PyAdapterContextFeatureProvider,
-                                                                  elementFeaturesProvider: PyAdapterElementFeatureProvider) -> Unit) {
-    val contextFeaturesProvider = PyAdapterContextFeatureProvider(PyContextFeatureProvider())
-    val elementFeaturesProvider = PyAdapterElementFeatureProvider(PyElementFeatureProvider())
-    try {
-      ContextFeatureProvider.EP_NAME.addExplicitExtension(PythonLanguage.INSTANCE, contextFeaturesProvider)
-      ElementFeatureProvider.EP_NAME.addExplicitExtension(PythonLanguage.INSTANCE, elementFeaturesProvider)
-      action(contextFeaturesProvider, elementFeaturesProvider)
-    }
-    finally {
-      ContextFeatureProvider.EP_NAME.removeExplicitExtension(PythonLanguage.INSTANCE, contextFeaturesProvider)
-      ElementFeatureProvider.EP_NAME.removeExplicitExtension(PythonLanguage.INSTANCE, elementFeaturesProvider)
-    }
-  }
-
-  private fun assertHasFeatures(actual: Map<String, MLFeatureValue>,
-                                expectedDefined: List<Pair<String, MLFeatureValue>>) {
-    for (pair in expectedDefined) {
-      Assert.assertTrue("Assert has feature: ${pair.first}", actual.containsKey(pair.first))
-      Assert.assertEquals("Check feature value: ${pair.first}", pair.second.toString(), actual[pair.first].toString())
-    }
-  }
-
-  private fun assertHasNotFeatures(actual: Map<String, MLFeatureValue>, expected: List<String>) {
-    for (value in expected) {
-      Assert.assertFalse("Assert has not feature: $value", actual.containsKey(value))
-    }
-  }
-
-  private fun invokeCompletion() {
-    myFixture.configureByFile(getTestName(true) + ".py")
-    myFixture.completeBasic()
-  }
-
-  private fun kwId(kw: String): Int {
-    return PyMlCompletionHelpers.getKeywordId(kw)!!
-  }
 }
@@ -0,0 +1,91 @@
+// Copyright 2000-2020 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.python.codeInsight.mlcompletion
+
+import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
+import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
+import com.intellij.codeInsight.lookup.LookupElement
+import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.fixtures.PyTestCase
+import org.junit.Assert
+
+open class PyMlCompletionTestCase: PyTestCase() {
+  fun doContextFeaturesTest(vararg expected: Pair<String, MLFeatureValue>) = doContextFeaturesTest(listOf(*expected), emptyList())
+
+  fun doContextFeaturesTest(expectedDefined: List<Pair<String, MLFeatureValue>>, expectedUndefined: List<String>) {
+    doWithInstalledProviders { contextFeaturesProvider, _ ->
+      invokeCompletion()
+      assertHasFeatures(contextFeaturesProvider.features, expectedDefined)
+      assertHasNotFeatures(contextFeaturesProvider.features, expectedUndefined)
+    }
+  }
+
+  fun doElementFeaturesTest(checks: List<Pair<String, List<Pair<String, MLFeatureValue>>>>) {
+    checks.forEach {
+      doElementFeaturesTest(it.first, it.second, emptyList())
+    }
+  }
+
+  fun doElementFeaturesTest(elementToSelect: String, vararg expected: Pair<String, MLFeatureValue>) {
+    doElementFeaturesTest(elementToSelect, arrayListOf(*expected), emptyList())
+  }
+
+  fun doElementFeaturesTest(elementToSelect: String,
+                                    expectedDefined: List<Pair<String, MLFeatureValue>>,
+                                    expectedUndefined: List<String>) {
+    val selector: (LookupElement) -> Boolean = { it.lookupString == elementToSelect }
+    doElementFeaturesInternalTest(selector, expectedDefined, expectedUndefined)
+  }
+
+  private fun doElementFeaturesInternalTest(selector: (LookupElement) -> Boolean,
+                                            expectedDefined: List<Pair<String, MLFeatureValue>>,
+                                            expectedUndefined: List<String>) {
+    doWithInstalledProviders { _, elementFeaturesProvider ->
+      invokeCompletion()
+
+      val selected = myFixture.lookupElements!!.find(selector)
+      assertNotNull(selected)
+
+      val features = elementFeaturesProvider.features[selected]
+      assertNotNull(features)
+      assertHasFeatures(features!!, expectedDefined)
+      assertHasNotFeatures(features, expectedUndefined)
+    }
+  }
+
+  fun doWithInstalledProviders(action: (contextFeaturesProvider: PyAdapterContextFeatureProvider,
+                                                elementFeaturesProvider: PyAdapterElementFeatureProvider) -> Unit) {
+    val contextFeaturesProvider = PyAdapterContextFeatureProvider(PyContextFeatureProvider())
+    val elementFeaturesProvider = PyAdapterElementFeatureProvider(PyElementFeatureProvider())
+    try {
+      ContextFeatureProvider.EP_NAME.addExplicitExtension(PythonLanguage.INSTANCE, contextFeaturesProvider)
+      ElementFeatureProvider.EP_NAME.addExplicitExtension(PythonLanguage.INSTANCE, elementFeaturesProvider)
+      action(contextFeaturesProvider, elementFeaturesProvider)
+    }
+    finally {
+      ContextFeatureProvider.EP_NAME.removeExplicitExtension(PythonLanguage.INSTANCE, contextFeaturesProvider)
+      ElementFeatureProvider.EP_NAME.removeExplicitExtension(PythonLanguage.INSTANCE, elementFeaturesProvider)
+    }
+  }
+
+  fun kwId(kw: String) = PyMlCompletionHelpers.getKeywordId(kw)!!
+
+  fun invokeCompletion() {
+    myFixture.configureByFile(getTestName(true) + ".py")
+    myFixture.completeBasic()
+  }
+
+  fun assertHasFeatures(actual: Map<String, MLFeatureValue>,
+                                expectedDefined: List<Pair<String, MLFeatureValue>>) {
+    for (pair in expectedDefined) {
+      Assert.assertTrue("Assert has feature: ${pair.first}", actual.containsKey(pair.first))
+      Assert.assertEquals("Check feature value: ${pair.first}", pair.second.toString(), actual[pair.first].toString())
+    }
+  }
+
+  fun assertHasNotFeatures(actual: Map<String, MLFeatureValue>, expected: List<String>) {
+    for (value in expected) {
+      Assert.assertFalse("Assert has not feature: $value", actual.containsKey(value))
+    }
+  }
+}
@@ -0,0 +1,86 @@
+// Copyright 2000-2020 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.python.codeInsight.mlcompletion
+
+import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PrevCallsModelsStorage
+
+class PyPrevCallsFeatureTest: PyMlCompletionTestCase() {
+  override fun getTestDataPath(): String = super.getTestDataPath() + "/codeInsight/mlcompletion/prev2calls"
+
+  fun testHaveOsModel() {
+    PrevCallsModelsStorage.haveModelForQualifier("os")
+    PrevCallsModelsStorage.haveModelForQualifier("os.path")
+  }
+
+  fun testOsPathPrimaryWeight() = checkFeaturesPresence("os", "path",
+                                                        listOf("prev_2_calls_primary_weight"),
+                                                        listOf("prev_2_calls_weight_empty_prev_calls",
+                                                               "prev_2_calls_weight_secondary_two_calls",
+                                                               "prev_2_calls_weight_one_call"))
+
+  fun testOsPathEmptyPrevCallsWeight() = checkFeaturesPresence("os", "path",
+                          listOf("prev_2_calls_weight_empty_prev_calls"),
+                          listOf("prev_2_calls_primary_weight",
+                                 "prev_2_calls_weight_secondary_two_calls",
+                                 "prev_2_calls_weight_one_call"))
+
+  fun testOsMakedirsPathPrimaryWeight() = checkFeaturesPresence("os", "path",
+                                                                listOf("prev_2_calls_primary_weight"),
+                                                                listOf("prev_2_calls_weight_one_call",
+                                                                      "prev_2_calls_weight_secondary_two_calls",
+                                                                      "prev_2_calls_weight_empty_prev_calls"))
+
+  fun testOsMakedirsPathOneCallWeight() = checkFeaturesPresence("os", "path",
+                                                                listOf("prev_2_calls_weight_one_call"),
+                                                                listOf("prev_2_calls_primary_weight",
+                                                                       "prev_2_calls_weight_secondary_two_calls",
+                                                                       "prev_2_calls_weight_empty_prev_calls"))
+
+  fun testOsEnvironGet() = checkFeaturesPresence("os.environ", "get",
+                                                 listOf("prev_2_calls_primary_weight"),
+                                                 listOf("prev_2_calls_weight_one_call",
+                                                                       "prev_2_calls_weight_secondary_two_calls",
+                                                                       "prev_2_calls_weight_empty_prev_calls"))
+
+  fun testSysStdinReadlineRstrip() = checkFeaturesPresence("sys.stdin.readline", "rstrip",
+                                                 listOf("prev_2_calls_primary_weight"),
+                                                 listOf("prev_2_calls_weight_one_call",
+                                                        "prev_2_calls_weight_secondary_two_calls",
+                                                        "prev_2_calls_weight_empty_prev_calls"))
+
+  fun testSysStdinReadlinesIter() = checkFeaturesPresence("sys.stdin", "readlines",
+                                                           listOf("prev_2_calls_primary_weight"),
+                                                           listOf("prev_2_calls_weight_one_call",
+                                                                  "prev_2_calls_weight_secondary_two_calls",
+                                                                  "prev_2_calls_weight_empty_prev_calls"))
+
+  fun testOsEnvironGetStartsWithCond() = checkFeaturesPresence("os.environ.get", "startswith",
+                                                               listOf("prev_2_calls_primary_weight"),
+                                                               listOf("prev_2_calls_weight_one_call",
+                                                                      "prev_2_calls_weight_secondary_two_calls",
+                                                                      "prev_2_calls_weight_empty_prev_calls"))
+
+  private fun checkFeaturesPresence(moduleName: String,
+                                    elementName: String,
+                                    expectedFeatures: List<String>,
+                                    notExpectedFeatures: List<String>) {
+    assertTrue(PrevCallsModelsStorage.haveModelForQualifier(moduleName))
+
+    doWithInstalledProviders { _, elementFeaturesProvider ->
+      invokeCompletion()
+
+      val selected = myFixture.lookupElements!!.find { it.lookupString == elementName }
+      assertNotNull(selected)
+
+      val features = elementFeaturesProvider.features[selected]
+      assertNotNull(features)
+
+      for (feature in expectedFeatures) {
+        assertTrue(features!!.containsKey(feature))
+      }
+
+      for (feature in notExpectedFeatures) {
+        assertFalse(features!!.containsKey(feature))
+      }
+    }
+  }
+}