PY-41056 Impl new ml completion features

GitOrigin-RevId: 2c7d1296854ced8f291aa7ee399ebd1881e2b005

plugins/stats-collector/src/com/intellij/completion/ml/MLFeaturesUtil.kt

@@ -4,6 +4,8 @@ package com.intellij.completion.ml
 import com.intellij.codeInsight.completion.CompletionLocation
 import com.intellij.codeInsight.completion.CompletionService
 import com.intellij.codeInsight.completion.CompletionSorter
+import com.google.common.cache.CacheBuilder
+import com.google.common.cache.CacheLoader
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementWeigher
@@ -21,9 +23,17 @@ object MLFeaturesUtil {
     }
   }
 
+  val classNameSafeCache = CacheBuilder
+    .newBuilder()
+    .softValues()
+    .maximumSize(100)
+    .build(object: CacheLoader<Class<*>, String>() {
+      override fun load(clazz: Class<*>) = if (getPluginInfo(clazz).isSafeToReport()) clazz.name else "third.party"
+    })
+
   private fun getClassNameSafe(feature: MLFeatureValue.ClassNameValue): String {
     val clazz = feature.value
-    return if (getPluginInfo(clazz).isSafeToReport()) clazz.name else "third.party"
+    return classNameSafeCache[clazz]
   }
 
   fun addWeighersToNonDefaultSorter(sorter: CompletionSorter, location: CompletionLocation, vararg weigherIds: String): CompletionSorter {

plugins/stats-collector/src/com/intellij/completion/ml/common/CommonElementLocationFeatures.kt

@@ -29,6 +29,10 @@ class CommonElementLocationFeatures : ElementFeatureProvider {
       }
     }
 
+    completionElement?.let {
+      result["item_class"] = MLFeatureValue.className(it::class.java)
+    }
+
     return result
   }
 

plugins/stats-collector/test/com/intellij/completion/ml/common/CommonElementLocationFeaturesTest.kt

@@ -9,6 +9,10 @@ import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.lang.Language
 import com.intellij.psi.PsiNamedElement
+import com.intellij.psi.impl.source.PsiFieldImpl
+import com.intellij.psi.impl.source.PsiMethodImpl
+import com.intellij.psi.impl.source.PsiParameterImpl
+import com.intellij.psi.impl.source.tree.java.PsiLocalVariableImpl
 import junit.framework.TestCase
 
 class CommonElementLocationFeaturesTest: LightCompletionTestCase() {
@@ -46,6 +50,25 @@ class CommonElementLocationFeaturesTest: LightCompletionTestCase() {
     assertFalse(features.containsKey("Exception"))
   }
 
+  fun `test lookup element psi class name`() {
+    val features = calculateFeature("lookup_element_psi_class_name", "Test.java") {
+      """|class Test {
+         |  private int a = 1;
+         |  
+         |  void f(String s) {
+         |    long c = 3;
+         |    System.out.println(<caret>);
+         |  }
+         |}
+      """.trimMargin()
+    }
+
+    assertFeaturesEquals(features.getValue("a"), MLFeatureValue.className(PsiFieldImpl::class.java))
+    assertFeaturesEquals(features.getValue("s"), MLFeatureValue.className(PsiParameterImpl::class.java))
+    assertFeaturesEquals(features.getValue("f"), MLFeatureValue.className(PsiMethodImpl::class.java))
+    assertFeaturesEquals(features.getValue("c"), MLFeatureValue.className(PsiLocalVariableImpl::class.java))
+  }
+
   @Suppress("SameParameterValue")
   private fun calculateFeature(featureName: String, fileName: String, text: () -> String): Map<String, MLFeatureValue> {
     val result: MutableMap<String, MLFeatureValue> = mutableMapOf()

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyCompletionFeatures.kt

@@ -2,7 +2,9 @@
 package com.jetbrains.python.codeInsight.mlcompletion
 
 import com.intellij.codeInsight.completion.CompletionLocation
+import com.intellij.codeInsight.completion.CompletionUtil
 import com.intellij.codeInsight.completion.ml.ContextFeatures
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementPresentation
 import com.intellij.openapi.module.ModuleUtilCore
@@ -24,10 +26,34 @@ object PyCompletionFeatures {
     return ("dict key" == presentation.typeText)
   }
 
-  fun isTheSameFile(element: LookupElement, location: CompletionLocation): Boolean {
-    val psiFile = location.completionParameters.originalFile
-    val elementPsiFile = element.psiElement?.containingFile ?: return false
-    return psiFile == elementPsiFile
+  fun getElementPsiLocationFeatures(element: LookupElement, location: CompletionLocation): Map<String, MLFeatureValue> {
+    val caretPsiPosition = location.completionParameters.position
+    val elementPsiPosition = element.psiElement ?: return emptyMap()
+
+    val caretFile = caretPsiPosition.containingFile?.originalFile ?: return emptyMap()
+    val elementFile = elementPsiPosition.containingFile?.originalFile ?: return emptyMap()
+    if (caretFile != elementFile) return emptyMap()
+
+    val result = mutableMapOf(
+      "is_the_same_file" to MLFeatureValue.binary(true),
+      "text_offset_distance" to MLFeatureValue.numerical(caretPsiPosition.textOffset - elementPsiPosition.textOffset)
+    )
+
+    if (isTheSameScope(caretPsiPosition, elementPsiPosition, PyClass::class.java)) {
+      result["is_the_same_class"] = MLFeatureValue.binary(true)
+    }
+
+    if (isTheSameScope(caretPsiPosition, elementPsiPosition, PyFunction::class.java)) {
+      result["is_the_same_method"] = MLFeatureValue.binary(true)
+    }
+
+    return result
+  }
+
+  private fun <T: PsiElement> isTheSameScope(caretPsiPosition: PsiElement, elementPsiPosition: PsiElement, scopeClass: Class<T>): Boolean {
+    val caretEnclosingScope = PsiTreeUtil.getParentOfType(caretPsiPosition, scopeClass) ?: return false
+    val elementEnclosingScope = PsiTreeUtil.getParentOfType(elementPsiPosition, scopeClass) ?: return false
+    return isOriginalElementsTheSame(caretEnclosingScope, elementEnclosingScope)
   }
 
   fun isTakesParameterSelf(element: LookupElement): Boolean {
@@ -186,13 +212,11 @@ object PyCompletionFeatures {
 
   fun getPyLookupElementInfo(element: LookupElement): PyCompletionMlElementInfo? = element.getUserData(PyCompletionMlElementInfo.key)
 
-  fun getNumberOfQualifiersInExpresionFeature(element: PsiElement): Int {
-    if (element !is PyQualifiedExpression) return 1
-    return element.asQualifiedName()?.components?.size ?: 1
-  }
-
   private fun isAfterColon(locationPsi: PsiElement): Boolean {
     val prevVisibleLeaf = PsiTreeUtil.prevVisibleLeaf(locationPsi)
     return (prevVisibleLeaf != null && prevVisibleLeaf.elementType == PyTokenTypes.COLON)
   }
+
+  private fun isOriginalElementsTheSame(a: PsiElement, b: PsiElement) =
+    CompletionUtil.getOriginalElement(a) == CompletionUtil.getOriginalElement(b)
 }

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyContextFeatureProvider.kt

@@ -6,6 +6,7 @@ import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PrevCallsModelsProviderService
 import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PyPrevCallsCompletionFeatures
+import com.jetbrains.python.psi.PyExpression
 import com.jetbrains.python.psi.types.TypeEvalContext
 
 class PyContextFeatureProvider : ContextFeatureProvider {
@@ -19,7 +20,11 @@ class PyContextFeatureProvider : ContextFeatureProvider {
     result["is_in_condition"] = MLFeatureValue.binary(PyCompletionFeatures.isInCondition(position))
     result["is_after_if_statement_without_else_branch"] = MLFeatureValue.binary(PyCompletionFeatures.isAfterIfStatementWithoutElseBranch(position))
     result["is_in_for_statement"] = MLFeatureValue.binary(PyCompletionFeatures.isInForStatement(position))
-    result["num_of_prev_qualifiers"] = MLFeatureValue.numerical(PyCompletionFeatures.getNumberOfQualifiersInExpresionFeature(position))
+
+    val positionParent = position.parent
+    if (positionParent is PyExpression) {
+      result["num_of_prev_qualifiers"] = MLFeatureValue.numerical(PyMlCompletionHelpers.getQualifiedComponents(positionParent).size)
+    }
 
     val neighboursKws = PyCompletionFeatures.getPrevNeighboursKeywordIds(position)
     if (neighboursKws.size > 0) result["prev_neighbour_keyword_1"] = MLFeatureValue.numerical(neighboursKws[0])
@@ -48,6 +53,7 @@ class PyContextFeatureProvider : ContextFeatureProvider {
     PyNamesMatchingMlCompletionFeatures.calculateNamedArgumentsNames(environment)
     PyNamesMatchingMlCompletionFeatures.calculateImportNames(environment)
     PyNamesMatchingMlCompletionFeatures.calculateStatementListNames(environment)
+    PyNamesMatchingMlCompletionFeatures.calculateEnclosingMethodName(environment)
 
     PyNamesMatchingMlCompletionFeatures.calculateSameLineLeftNames(environment).let { names ->
       result["have_opening_round_bracket"] = MLFeatureValue.binary(PyParenthesesFeatures.haveOpeningRoundBracket(names))

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyElementFeatureProvider.kt

@@ -7,6 +7,7 @@ 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
+import com.jetbrains.python.psi.PyParameter
 
 class PyElementFeatureProvider : ElementFeatureProvider {
   override fun getName(): String = "python"
@@ -18,6 +19,7 @@ class PyElementFeatureProvider : ElementFeatureProvider {
 
     val lookupString = element.lookupString
     val locationPsi = location.completionParameters.position
+    val lookupPsiElement = element.psiElement
 
     PyCompletionFeatures.getPyLookupElementInfo(element)?.let { info ->
       result["kind"] = MLFeatureValue.categorical(info.kind)
@@ -35,13 +37,14 @@ class PyElementFeatureProvider : ElementFeatureProvider {
     }
 
     result["is_dict_key"] = MLFeatureValue.binary(PyCompletionFeatures.isDictKey(element))
-    result["is_the_same_file"] = MLFeatureValue.binary(PyCompletionFeatures.isTheSameFile(element, location))
     result["is_takes_parameter_self"] = MLFeatureValue.binary(PyCompletionFeatures.isTakesParameterSelf(element))
     result["underscore_type"] = MLFeatureValue.categorical(PyCompletionFeatures.getElementNameUnderscoreType(lookupString))
     result["number_of_tokens"] = MLFeatureValue.numerical(PyNamesMatchingMlCompletionFeatures.getNumTokensFeature(lookupString))
     result["element_is_py_file"] = MLFeatureValue.binary(PyCompletionFeatures.isPsiElementIsPyFile(element))
     result["element_is_psi_directory"] = MLFeatureValue.binary(PyCompletionFeatures.isPsiElementIsPsiDirectory(element))
 
+    result.putAll(PyCompletionFeatures.getElementPsiLocationFeatures(element, location))
+
     PyCompletionFeatures.getElementModuleCompletionFeatures(element)?.let { with(it) {
         result["element_module_is_std_lib"] = MLFeatureValue.binary(isFromStdLib)
         result["can_find_element_module"] = MLFeatureValue.binary(canFindModule)
@@ -78,6 +81,14 @@ class PyElementFeatureProvider : ElementFeatureProvider {
       result["receiver_tokens_num"] = MLFeatureValue.numerical(receiverTokensNum)
     }}
 
+    result.putAll(PyNamesMatchingMlCompletionFeatures.getMatchingWithEnclosingMethodFeatures(contextFeatures, element))
+
+    if (lookupString.endsWith("Warning")) result["is_warning"] = MLFeatureValue.binary(true)
+    if (lookupString.endsWith("Error")) result["is_error"] = MLFeatureValue.binary(true)
+    if (lookupString.endsWith("Exception")) result["is_exception"] = MLFeatureValue.binary(true)
+    if (lookupString.endsWith("s")) result["ends_with_s"] = MLFeatureValue.binary(true)
+    if (lookupPsiElement is PyParameter && lookupPsiElement.isSelf) result["is_self"] = MLFeatureValue.binary(true)
+
     contextFeatures.getUserData(PyPrevCallsCompletionFeatures.PREV_CALLS_CONTEXT_INFO_KEY)?.let { contextInfo ->
       PyPrevCallsCompletionFeatures.getResult(lookupString, contextInfo)?.let { with (it) {
         when {

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyMlCompletionHelpers.kt

@@ -4,6 +4,9 @@ package com.jetbrains.python.codeInsight.mlcompletion
 import com.google.common.reflect.TypeToken
 import com.google.gson.Gson
 import com.intellij.openapi.diagnostic.Logger
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.psi.PyExpression
+import com.jetbrains.python.psi.PyQualifiedExpression
 import java.io.InputStreamReader
 import java.nio.charset.StandardCharsets
 
@@ -30,4 +33,11 @@ object PyMlCompletionHelpers {
       return emptyMap()
     }
   }
+
+  fun getQualifiedComponents(element: PyExpression): List<String> =
+    generateSequence<PsiElement>(element) { it.firstChild }
+      .filterIsInstance<PyQualifiedExpression>()
+      .mapNotNull { it.name }
+      .toList()
+      .asReversed()
 }

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyNamesMatchingMlCompletionFeatures.kt

@@ -4,6 +4,7 @@ package com.jetbrains.python.codeInsight.mlcompletion
 import com.intellij.codeInsight.completion.CompletionUtilCore.DUMMY_IDENTIFIER_TRIMMED
 import com.intellij.codeInsight.completion.ml.CompletionEnvironment
 import com.intellij.codeInsight.completion.ml.ContextFeatures
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.text.StringUtil
@@ -26,6 +27,8 @@ object PyNamesMatchingMlCompletionFeatures {
   val statementListOrFileNamesKey = Key<Map<String, Int>>("py.ml.completion.statement.list.names")
   val statementListOrFileTokensKey = Key<Map<String, Int>>("py.ml.completion.statement.list.tokens")
 
+  private val enclosingMethodName = Key<String>("py.ml.completion.enclosing.method.name")
+
   data class PyScopeMatchingFeatures(val sumMatches: Int,
                                      val sumTokensMatches: Int,
                                      val numScopeNames: Int,
@@ -57,6 +60,14 @@ object PyNamesMatchingMlCompletionFeatures {
     return MatchingWithReceiverFeatures(matchesWithReceiver, receiverTokensNum, numMatchedTokens)
   }
 
+  fun getMatchingWithEnclosingMethodFeatures(contextFeatures: ContextFeatures, element: LookupElement): Map<String, MLFeatureValue> {
+    val name = contextFeatures.getUserData(enclosingMethodName) ?: return emptyMap()
+    val result = mutableMapOf<String, MLFeatureValue>()
+    if (element.lookupString == name) result["matches_with_enclosing_method"] = MLFeatureValue.binary(true)
+    result["matched_tokens_with_enclosing_method"] = MLFeatureValue.numerical(tokensMatched (name, element.lookupString))
+    return result
+  }
+
   fun calculateFunBodyNames(environment: CompletionEnvironment) {
     val position = environment.parameters.position
     val scope = PsiTreeUtil.getParentOfType(position, PyFile::class.java, PyFunction::class.java, PyClass::class.java)
@@ -108,6 +119,12 @@ object PyNamesMatchingMlCompletionFeatures {
       ?.let { putTokensAndNamesToUserData(environment, namedArgumentsNamesKey, namedArgumentsTokensKey, it) }
   }
 
+  fun calculateEnclosingMethodName(environment: CompletionEnvironment) {
+    val position = environment.parameters.position
+    val name = PsiTreeUtil.getParentOfType(position, PyFunction::class.java)?.name ?: return
+    environment.putUserData(enclosingMethodName, name)
+  }
+
   private fun putTokensAndNamesToUserData(environment: CompletionEnvironment,
                                           namesKey: Key<Map<String, Int>>,
                                           tokensKey: Key<Map<String, Int>>,

python/src/com/jetbrains/python/codeInsight/mlcompletion/prev2calls/AssignmentVisitor.kt

@@ -2,6 +2,7 @@
 package com.jetbrains.python.codeInsight.mlcompletion.prev2calls
 
 import com.intellij.psi.PsiElement
+import com.jetbrains.python.codeInsight.mlcompletion.PyMlCompletionHelpers
 import com.jetbrains.python.psi.*
 
 class AssignmentVisitor(private val borderOffset: Int,
@@ -39,7 +40,7 @@ class AssignmentVisitor(private val borderOffset: Int,
       val left = it.first
       val right = it.second
       if (left is PyTargetExpression) {
-        val leftName = getQualifierComponents(left).joinToString(".")
+        val leftName = PyMlCompletionHelpers.getQualifiedComponents(left).joinToString(".")
         val rightName = getResolvedExpression(right).resolvedExpression
         if (rightName.isNotEmpty() && leftName.isNotEmpty()) {
           fullNames[leftName] = rightName
@@ -49,10 +50,10 @@ class AssignmentVisitor(private val borderOffset: Int,
   }
 
   data class ResolvedExpression(val resolvedExpression: String = "", val resolvedPrefix: String = "")
-  private fun getResolvedExpression(node: PsiElement?): ResolvedExpression {
+  private fun getResolvedExpression(node: PyExpression?): ResolvedExpression {
     if (node == null) return ResolvedExpression()
 
-    val components = getQualifierComponents(node)
+    val components = PyMlCompletionHelpers.getQualifiedComponents(node)
     for (i in components.indices) {
       val firstN = i + 1
       val prefix = components.take(firstN).joinToString(".")
@@ -68,13 +69,6 @@ class AssignmentVisitor(private val borderOffset: Int,
     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)
   }

python/testData/codeInsight/mlcompletion/locationDifferentFile.py

@@ -0,0 +1 @@
+<caret>

python/testData/codeInsight/mlcompletion/locationSameFileAndClass.py

@@ -0,0 +1,6 @@
+class Clazzz:
+  def some_function(self):
+    self.some_variable = 42
+
+  def foo(self):
+    self.<caret>

python/testData/codeInsight/mlcompletion/locationSameFileAndMethod.py

@@ -0,0 +1,3 @@
+def foo():
+  some_variable = 42
+  <caret>

python/testData/codeInsight/mlcompletion/locationSameFileAndMethodAndClass.py

@@ -0,0 +1,4 @@
+class Clazzz:
+  def foo(self):
+    some_variable = 42
+    <caret>

python/testData/codeInsight/mlcompletion/locationSameFileOnly.py

@@ -0,0 +1,4 @@
+def some_function():
+  print(42)
+
+<caret>

python/testData/codeInsight/mlcompletion/matchesWithEnclosingMethodAlmostTheSameName.py

@@ -0,0 +1,3 @@
+_qwer_tyuio_asdf_gh = 123
+def qwer_tyuio_asdf_gh():
+  print(<caret>)

python/testData/codeInsight/mlcompletion/matchesWithEnclosingMethodTheSameName.py

@@ -0,0 +1,3 @@
+_qwer_tyuio_asdf_gh = 123
+def _qwer_tyuio_asdf_gh():
+  print(<caret>)

python/testData/codeInsight/mlcompletion/numOfPrevQualifiersIs1.py

@@ -0,0 +1 @@
+a<caret>

python/testData/codeInsight/mlcompletion/numOfPrevQualifiersIs3.py

@@ -0,0 +1,5 @@
+class Clzz:
+  pass
+qwer = Clzz()
+qwer._asdFjk_ = "asd"
+qwer._asdFjk_.<caret>

python/testData/codeInsight/mlcompletion/numOfPrevQualifiersIs4.py

@@ -0,0 +1,10 @@
+class Clzz:
+  def b(self):
+    return Clzz()
+
+  def c(self):
+    return "123"
+
+
+a = Clzz()
+a.b().c().<caret>

python/testSrc/com/jetbrains/python/codeInsight/mlcompletion/PyMlCompletionFeaturesTest.kt

@@ -3,6 +3,8 @@ package com.jetbrains.python.codeInsight.mlcompletion
 
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.impl.PyFunctionImpl
+import com.jetbrains.python.psi.impl.PyTargetExpressionImpl
 
 class PyMlCompletionFeaturesTest: PyMlCompletionTestCase() {
   override fun getTestDataPath(): String = super.getTestDataPath() + "/codeInsight/mlcompletion"
@@ -100,6 +102,12 @@ class PyMlCompletionFeaturesTest: PyMlCompletionTestCase() {
   fun testInsideClassAfterConstructor() = doContextFeaturesTest(Pair("containing_class_have_constructor", MLFeatureValue.binary(true)),
                                                                 Pair("diff_lines_with_class_def", MLFeatureValue.numerical(4)))
 
+  fun testNumOfPrevQualifiersIs3() = doContextFeaturesTest(Pair("num_of_prev_qualifiers", MLFeatureValue.numerical(3)))
+
+  fun testNumOfPrevQualifiersIs4() = doContextFeaturesTest(Pair("num_of_prev_qualifiers", MLFeatureValue.numerical(4)))
+
+  fun testNumOfPrevQualifiersIs1() = doContextFeaturesTest(Pair("num_of_prev_qualifiers", MLFeatureValue.numerical(1)))
+
   // Element features
 
   fun testDictKey() = doElementFeaturesTest("\"dict_key\"",
@@ -245,4 +253,49 @@ class PyMlCompletionFeaturesTest: PyMlCompletionTestCase() {
                                                           Pair("receiver_name_matches", MLFeatureValue.binary(false)),
                                                           Pair("receiver_num_matched_tokens", MLFeatureValue.numerical(3)),
                                                           Pair("receiver_tokens_num", MLFeatureValue.numerical(3)))
+
+  fun testMatchesWithEnclosingMethodTheSameName() = doElementFeaturesTest(
+    "_qwer_tyuio_asdf_gh",
+    Pair("number_of_tokens", MLFeatureValue.numerical(4)),
+    Pair("matches_with_enclosing_method", MLFeatureValue.binary(true)),
+    Pair("matched_tokens_with_enclosing_method", MLFeatureValue.numerical(4))
+  )
+
+  fun testMatchesWithEnclosingMethodAlmostTheSameName() = doElementFeaturesTest(
+    "_qwer_tyuio_asdf_gh",
+    listOf(Pair("number_of_tokens", MLFeatureValue.numerical(4)),
+    Pair("matched_tokens_with_enclosing_method", MLFeatureValue.numerical(4))),
+    listOf("matches_with_enclosing_method")
+  )
+
+  fun testLocationSameFileAndMethodAndClass() = doElementFeaturesTest(
+    "some_variable",
+    Pair("is_the_same_file", MLFeatureValue.binary(true)),
+    Pair("is_the_same_class", MLFeatureValue.binary(true)),
+    Pair("is_the_same_method", MLFeatureValue.binary(true))
+  )
+
+  fun testLocationSameFileAndClass() = doElementFeaturesTest(
+    "some_variable",
+    listOf(Pair("is_the_same_file", MLFeatureValue.binary(true)), Pair("is_the_same_class", MLFeatureValue.binary(true))),
+    listOf("is_the_same_method")
+  )
+
+  fun testLocationSameFileAndMethod() = doElementFeaturesTest(
+    "some_variable",
+    listOf(Pair("is_the_same_file", MLFeatureValue.binary(true)), Pair("is_the_same_method", MLFeatureValue.binary(true))),
+    listOf("is_the_same_class")
+  )
+
+  fun testLocationSameFileOnly() = doElementFeaturesTest(
+    "some_function",
+    listOf(Pair("is_the_same_file", MLFeatureValue.binary(true))),
+    listOf("is_the_same_class", "is_the_same_method")
+  )
+
+  fun testLocationDifferentFile() = doElementFeaturesTest(
+    "min",
+    emptyList(),
+    listOf("is_the_same_file", "is_the_same_class", "is_the_same_method")
+  )
 }