如何为Kotlin类添加带另一接收者类型的扩展函数？

你想要的这种仅允许在特定类的上下文里调用另一个类的扩展函数的需求，正好可以用Kotlin的**上下文接收者（Context Receivers）**特性来实现，无需修改第三方库的源码、创建子类或伴生对象。

核心实现思路 #

上下文接收者允许函数拥有多个接收者：一个是函数的直接接收者（比如示例中的B），另一个是限定调用上下文的接收者（比如示例中的A）。只有当代码处于限定接收者的上下文范围内时，才能调用这个扩展函数。

针对你给出的A/B简化示例的实现 #

/** 第三方库中的类，不可修改 **/
abstract class A {
    fun B.doSomeThing() {
        println("something")
    }
}

class B

/*********************************************/

// 用上下文接收者定义多接收者扩展函数
context(A)
fun B.doSomeThingElse() {
    println("somethingElse")
}

// 测试代码
class AImpl: A() {
    val b = B()
    fun main() {
        b.doSomeThing() // 正常运行
        b.doSomeThingElse() // 正常运行：当前处于A的上下文
    }
}

class C {
    val b = B()
    fun main() {
        b.doSomeThingElse() // 编译错误：无A的上下文，无法调用
    }
}

对应ViewModel+Flow场景的实际实现 #

把这个思路迁移到你最初的需求上，实现Flow的扩展函数仅能在ViewModel上下文调用：

context(ViewModel)
fun <T> Flow<T>.launchInViewModelScope(block: suspend (T) -> Unit) {
    // 直接访问ViewModel的viewModelScope
    this.onEach(block).launchIn(viewModelScope)
}

在ViewModel子类中调用时，直接使用即可：

class MyViewModel : ViewModel() {
    val myFlow = flow { emit("test") }

    fun init() {
        myFlow.launchInViewModelScope { data ->
            // 处理Flow数据
        }
    }
}

注意事项 #

1. Kotlin版本要求：上下文接收者在Kotlin 1.7及以上正式支持；若使用Kotlin 1.6，需开启实验性支持，在Gradle配置中添加：

   // Gradle Kotlin DSL
   kotlinOptions {
       freeCompilerArgs += "-Xcontext-receivers"
   }
   

   或Groovy语法：

   // Gradle Groovy
   kotlinOptions {
       freeCompilerArgs += '-Xcontext-receivers'
   }
   

2. 上下文接收者会严格限定函数的调用范围，确保只有在ViewModel（或示例中的A）的实例上下文里才能调用目标扩展函数，完全符合你的需求。

内容的提问来源于stack exchange，提问作者lolo.io