keyを指定したViewModelの取得もlazy拡張関数で簡単にしたい

ViewModelを取得するとき、activity-ktxやfragment-ktxで提供されている拡張関数を使うと、ViewModelProviderなどの記述を省略でき、すっきりと記述できます。

HogeActivity.kt

private val viewModel: HogeViewModel by viewModels()

HogeFragment.kt

private val viewModel: HogeViewModel by viewModels()
private val viewModel: HogeViewModel by activityViewModels()

これは大変便利なのですが、keyを指定することができません。
例えば、ViewPager上のFragmentで使用するViewModelなどはActivityのViewModelStoreを使い、各ページをキーとしたViewModelを使う必要がありますが、この場合には使えないことになります。

まあ、使えなかったところでそこまで複雑な記述になるわけでもないのですが

HogeFragment.kt

private val viewModel: PageViewModel by lazy { 
     ViewModelProvider(requireActivity()).get("key", PageViewModel::class.java)
}

ViewModelをkey指定で取得するときはkeyの名前空間に注意しようで書いたように、keyの名前空間への配慮も必要なので、まとめて面倒見てくれる拡張関数が欲しくなります。

既存の仕組みでできないのか

ではやり方を調べるため、まずは引数の追加とかでできないか、ソースを追ってみましょう。

ActivityViewModelLazy.kt

@MainThread
inline fun <reified VM : ViewModel> ComponentActivity.viewModels(
    noinline factoryProducer: (() -> Factory)? = null
): Lazy<VM> {
    val factoryPromise = factoryProducer ?: {
        defaultViewModelProviderFactory
    }

    return ViewModelLazy(VM::class, { viewModelStore }, factoryPromise)
}

ViewModelProvider.kt

public class ViewModelLazy<VM : ViewModel> (
    private val viewModelClass: KClass<VM>,
    private val storeProducer: () -> ViewModelStore,
    private val factoryProducer: () -> ViewModelProvider.Factory
) : Lazy<VM> {
    private var cached: VM? = null

    override val value: VM
        get() {
            val viewModel = cached
            return if (viewModel == null) {
                val factory = factoryProducer()
                val store = storeProducer()
                ViewModelProvider(store, factory).get(viewModelClass.java).also {
                    cached = it
                }
            } else {
                viewModel
            }
        }

    override fun isInitialized(): Boolean = cached != null
}

以上！シンプルですね。ViewModelLazyでViewModelProviderを呼び出していますが、keyを指定する余地がありません。

ないのなら作ってしまおう

やっていることは非常にシンプルなので、このViewModelLazyにkeyを指定できるようにしたクラスを用意して、同様にkeyを渡せる拡張関数を作れば良さそうです。

ViewModelExtensions.kt

@MainThread
inline fun <reified VM : ViewModel> ComponentActivity.keyedViewModels(
    noinline keyProducer: () -> String,
    noinline factoryProducer: (() -> ViewModelProvider.Factory)? = null
): Lazy<VM> = KeyedViewModelLazy(
    VM::class,
    keyProducer,
    { viewModelStore },
    factoryProducer ?: { defaultViewModelProviderFactory }
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.keyedViewModels(
    noinline keyProducer: () -> String,
    noinline ownerProducer: () -> ViewModelStoreOwner = { this },
    noinline factoryProducer: (() -> ViewModelProvider.Factory)? = null
): Lazy<VM> = KeyedViewModelLazy(
    VM::class,
    keyProducer,
    { ownerProducer().viewModelStore },
    factoryProducer ?: { defaultViewModelProviderFactory }
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.keyedActivityViewModels(
    noinline keyProducer: () -> String,
    noinline factoryProducer: (() -> ViewModelProvider.Factory)? = null
): Lazy<VM> = KeyedViewModelLazy(
    VM::class,
    keyProducer,
    { requireActivity().viewModelStore },
    factoryProducer ?: { requireActivity().defaultViewModelProviderFactory }
)

class KeyedViewModelLazy<VM : ViewModel>(
    private val viewModelClass: KClass<VM>,
    private val keyProducer: () -> String,
    private val storeProducer: () -> ViewModelStore,
    private val factoryProducer: () -> ViewModelProvider.Factory
) : Lazy<VM> {
    private var cached: VM? = null

    override val value: VM
        get() = cached
            ?: ViewModelProvider(storeProducer(), factoryProducer())
                .get(
                    viewModelClass.qualifiedName + ":" + keyProducer(),
                    viewModelClass.java
                )
                .also { cached = it }

    override fun isInitialized(): Boolean = cached != null
}

こんなところでしょうか、keyにはモデルクラス名をプレフィックスとしてつけるようにしているので、keyの名前空間についても気にする必要は無くなります。

既存のviewModels``activityViewModelsと名前がかぶるのはよろしくないのでkeyedをつけています。その関係でkey指定なしには対応していません。なんだかどこかですでにありそうですが。

以上です。