How to Set up Dagger With ViewModel & Saved State Module

Android Developers finally @Provide an opinion about dependency injection, and the winner is... Dagger!

From my perspective, Dagger is “not great, not terrible." But it scales well and, when properly set up, you don't need to “touch” it later.

In this article, I try to explain how to set up Dagger to work with ViewModel and SavedState module in the most universal way — set it once, use it forever.

WARNING:
I’m not going to get deep into Dagger, so if you don’t know the basics of how Dagger works, please check it out first. Otherwise, the article may be tricky to understand.

TLDR:
I'm sorry, this article is longer than I expected. But I wanted to explain everything with a sufficient amount of information. For those who are familiar with the details and prefer seeing working snippets only (or you're just a little lazy), skip ahead to the TLDR section, at the bottom of this page.

Fasten your seatbelt! It's gonna be a rough ride.

1_fK-6zA16gBRDnK1ySl8vzg-1
Source: Giphy, https://gph.is/2Lhtrua

THE MOTIVATION

Let's start with motivation: Why do we want to do all of this?

Say you have a ViewModel class SomeViewModel : ViewModel(). Now, you want to fully use the power of ViewModels, so you apply inversion of control and pass dependencies to the constructor.

class SomeViewModel(
  private val dep1: Dependency,
  private val dep2: Dependency2
) : ViewModel()

If you're even more demanding, you have DI framework (Dagger in our case) do this for you. I won't describe how to set up Dagger with ViewModel, as there are many articles and SO answers available (e.g. here and here).

class SomeViewModel @Inject constructor(
  private val depFromDagger1: Dependency,
  private val depFromDagger2: Dependency2
) : ViewModel()

Still want more? Outrageous! Say you actually want to pass something besides dependencies coming from DI graph. Something like Bundle, or just some variable like articleId.

What are your options?

You can instantiate your ViewModel with dependencies from the graph and set your custom variable manually after construction with either lateinit var or with var of nullable type.

class SomeViewModel @Inject constructor(
  private val depFromDagger1: Dependency,
  private val depFromDagger2: Dependency2
) : ViewModel(){
  
  // may crash with UninitializedPropertyAccessException
  lateinit var articleId: String
  
  // must !! or ?. for every access
  var fragmentParams: Bundle? = null
}

Neither one of these is great because they makes your code more fragile. What if you reuse your ViewModel in another screen and forget to set the dynamic parameters? Crash! ...Or an improperly initialized class. And in case of a nullable variable, either you force !! it, or you have unnecessary null checks in your code.

Another option is to leave Dagger out of the game, create a custom ViewModel factory and manually pass dependencies which will be injected into Fragment or Activity. Not great.

class SomeViewModel(
  private val dep1: Dependency,
  private val dep2: Dependency2,
  private val articleId: String 
) : ViewModel() {
  
class Factory constructor(
    private val dep1: Dependency,
    private val dep2: Dependency2,
    private val articleId: String 
  ) : ViewModelProvider.Factory {
  
    override fun <T : ViewModel?> create(modelClass: Class<T>): T {
      return SomeViewModel(dep1, dep2, articleId) as T
    }
  } 
}

In Fragment, you need to @Inject the dependencies and pass it to the factory.

class SomeFragment : Fragment() {
  
  // inject into this class omitted for brevity
  
  @Inject
  lateinit var dep1: Dependency
  @Inject
  lateinit var dep2: Dependency
  
  lateinit var viewModel : SomeViewModel
  
  override fun onCreate(savedState: Bundle?){
    super.onCreate(savedState)
    // retrieve articleId and pass it to your factory
    val articleId = arguments!!.getString("article_id") 
    val factory = SomeViewModel.Factory(dep1, dep2, articleId)

    viewModel = ViewModelProvider(this, factory).get(SomeViewModel::class.java)
  }
}

This works, but it’s so much boilerplate. With each added, changed or removed dependency, you have to update three places in your code:

  • ViewModel’s constructor
  • ViewModel's custom factory
  • Instantiation of the factory with injected dependencies

(Un)fortunately, this is (probably) exactly what we want to achieve without all of the boilerplate, because we want to use one dynamically retrieved parameter — SavedStateHandle, from Saved State Module library.

VIEWMODEL AND ONSAVEINSTANCESTATE()

Before we dive into Saved State module, let's recap ViewModel's strengths and weaknesses.

ViewModel is great at handling orientation changes, as it survives when Fragment or Activity is destroyed. This lets you keep doing any long action without leaking the screen or needing to “restart” the action.

On the other hand, when you put your app into the background and the Android system kills it (usually due to longer inactivity and a need for more resources), the state of ViewModel is not preserved. What's worse, Fragment or Activity will handle this use case by calling onSaveInstanceState(outState: Bundle), but your ViewModel has no information about it. You need to handle this yourself by taking data from ViewModel and saving it in Fragment/Activity and later restoring it by manually placing it into ViewModel. But as I showed you earlier, this makes your ViewModel less robust because you can't run actions from constructor.

ViewModel doesn’t handle saving/restoring state.


The truth is, many apps don’t even bother solving this issue, which leads to weird behavior or even crashing when the user opens the app after some inactivity.

However, there's light at the end of the tunnel.

SAVED STATE MODULE FOR VIEWMODEL

Saved State Module for ViewModel is the new AndroidX library that allows handling instance state from ViewModel without any difficulty. This library provides custom factory for creating ViewModels. ViewModel constructor then expects SavedStateHandle parameter, which it communicates with. Destruction (and saving instance state) of a Fragment/Activity is then reflected in the handle, and ViewModel can therefore save or restore its state without the help of other classes.

SO HOW TO USE IT?

0.

Add gradle dependency.

implementation 
  "androidx.lifecycle:lifecycle-viewmodel-savedstate:1.0.0-rc03"

1.

Get the ViewModel with SavedStateViewModelFactory factory. You may also specify some default arguments to be passed into ViewModel. This parameter in the constructor is SavedStateRegistryOwner, which is either Fragment or Activity and serves as reference for saving/restoring the instance state.

class SomeFragment : Fragment() {
  
  lateinit var viewModel : SomeViewModel
  
  override fun onCreate(savedState: Bundle?){
    super.onCreate(savedState)    
    // default arguments, so you can set something dynamically
    val defaultArgs: Bundle? = bundleOf("id" to 5) // may be null
    
    // default factory for ViewModel creation
    val factory = SavedStateViewModelFactory(application, this, defaultArgs)
    
    // get the ViewModel with the factory and scope you want
    viewModel = ViewModelProvider(this, factory)[SomeViewModel::class.java]
  }
  
}

2.

In your ViewModel’s constructor, have variable of typeSavedStateHandle, which serves as a handler for saving or retrieving data. If you passed any default arguments, they will be part of this handle.

class SomeViewModel(
   private val application: Application,
   private val savedStateHandle: SavedStateHandle
){
  // ...
}

The handle has set/get methods similar to Bundle which, in case of ViewModel being killed by the system, are safely stored in app's state and restored later. It also has getLiveData(key) method, which returns MutableLiveData to simplify working with UI.

In the example below, counter value is retrieved in ViewModel's init{}. If value is not set, null is returned. Method onPlusClick() changes counter LiveData. In the end, you observe the LiveData and set current value into the handle, which will be safely stored.

class SomeViewModel(/* ..ommited.. */){
  val counter = MutableLiveData<Int>(0)

  init {
    counter.value = savedStateHandle.get("counter") ?: 0
    
    counter.observeForever { newValue ->
      savedStateHandle.set("counter", newValue)
    }
  }

  fun onPlusClick() {
    counter.value = (counter.value ?: 0) + 1
  }
}

You can also simplify it by using getLiveData(key, defaultValue) method, which results in the same functionality with less code.

class SomeViewModel(/* ..ommited.. */){
  val counter = savedStateHandle.getLiveData("counter", 0)

  fun onPlusClick() {
    counter.value = (counter.value ?: 0) + 1
  }
}

This way, when system kills your app, you can be sure that ViewModel will save and restore its state properly without the help of Fragment/Activity class.

HOW TO TEST THE SYSTEM KILLING YOUR APP?

With all this information, you're probably thinking: How do I know if it actually works? You often don’t want to sit and wait until the Android system kills your app in the background; you want to test it reasonably. There are two easy options that let you simulate the behavior:

a. Set the background process limit to none:
On your device, go to SettingsDeveloper options ➡ Background process limit set to No background processes

or

b. Kill your app with adb:
- Put your app into background (if you don’t do this step, adb will kill the app anyway, but the system won’t save the instance state).
- Use the adb command from snippet to kill your package. Small shell script may help to kill it for you and save some time.

#!/bin/bash

# Provide package of your application (com.example.myapp)
PACKAGE=$1

# First, put your app to background and then run this script

echo "Killing $PACKAGE"

adb shell ps | grep $PACKAGE | awk '{print $2}' | xargs adb shell run-as $PACKAGE kill

Alright, so — in the first section, we saw how to connect Dagger • ViewModel dot. This section described how to connect ViewModel • SavedStateHandle dot.
The question now is, how do we connect all the dots? How do we use Dagger and ViewModel with SavedStateHandle?

DAGGER AND VIEWMODEL WITH SAVEDSTATEHANDLE

As we know now, for ViewModel we can either have custom Factory which manually passes parameters, or fully Dagger-controlled instantiation. In our case, we would like to have a bit of both worlds.

Let's go over what sorcery needs to be done to connect the dots.

@AssistedInject for the win

Unfortunately, we cannot do it without some help—a hack, so to speak. We need to introduce a way to keep injecting Dagger-related dependencies (a Repository or application Context or something similar), while also having some of the parameters injected manually.

See how AssistedInject works here.

Fortunately, this is exactly what AssistedInject library does. It allows you to annotate the constructor of your class with @AssistedInject (in our case, ViewModel) and your dynamic parameters with @Assisted(in our case, SavedStateHandle). Based on this information, the library generates Factory compatible with Dagger for instantiation of your class. The factory has @Inject constructor with all dependencies from your ViewModel constructor and one method with the dynamic parameters marked with @Assisted.

(You can listen to Jake Wharton's talk, Helping Dagger Help You, where he describes it more in depth.)

HOW TO SET UP ASSISTEDINJECT WITH SAVEDSTATEHANDLE?

Let's start with the setup.

0.

Add gradle dependency.

compileOnly 
  'com.squareup.inject:assisted-inject-annotations-dagger2:0.5.2'
  
kapt 
  'com.squareup.inject:assisted-inject-processor-dagger2:0.5.2'

Version 0.5.2 indicates, that it may not be extra mature, but for our use case it works fine. And since the library only generates Dagger code, you shouldn’t see unstable code during execution of the app (worst case scenario, there may be issues with building).

1.

Create base interface which has one create(savedStateHandle: SavedStateHandle): T method. The generic parameter T serves as a way to inform AssistedInject about what type of ViewModel it should generate. This common interface allows us to omit some repetitive code in each ViewModel, and it's needed because we will then add all of our assisted factories to one @Multibinds map (shown later in this article).

/**
 * Base interface for all ViewModel factories
 */
interface AssistedSavedStateViewModelFactory<T : ViewModel> {
   fun create(savedStateHandle: SavedStateHandle): T
}

2.

Instead of @Inject, use @AssistedInject annotation. SavedStateHandle parameter annotate with @Assisted. This way the library knows which parameters are provided via Dagger and which are provided dynamically.

class SomeViewModel @AssistedInject constructor(
   private val application: Application,
   @Assisted private val savedStateHandle: SavedStateHandle
){
  // ...
}

3.

Inside of the ViewModel class, add interface annotated with @AssistedInject.Factory , which extends the base interface specified earlier — AssistedSavedStateViewModelFactory<T>.

class SomeViewModel @AssistedInject constructor(
    private val application: Application,
    @Assisted private val savedStateHandle: SavedStateHandle
){
    // must be inside of the ViewModel class!
    @AssistedInject.Factory
    interface Factory : AssistedSavedStateViewModelFactory<SomeViewModel> {
        override fun create(savedStateHandle: SavedStateHandle): SomeViewModel  // may be ommited prior kotlin 1.3.60 or after PR #121 in AssistedInject lib
    }
}

AssistedInject will generate class SomeViewModel_AssistedFactory, which is implementation of the Factory interface. It will have Dagger-related variables (Provider<*>) in its constructor and implementation of the create(savedStateHandle: SavedStateHandle) method from the base interface.

Side note: In Kotlin version 1.3.60 and above, you have to override the create method because there's a bug — super class uses the wrong name of argument, which AssistedInject doesn't expect. You can either override the method—as shown below—or check the issue in the library . There's also PR to handle different names of parameters, which solves the issue as well.

4.

Create Dagger @Module, mark it with @AssistedModuleand make it include AssistedInject_NameOfTheModule::class. The included module will be generated for you by the library. You have to build the project twice, so you can import it (same as with DaggerAppComponent being generated from AppComponent).

This generated module (AssistedInject_NameOfTheModule::class) contains @Binds methods from the generated ViewModel factories to its implementations (so that you can later inject your interface).

Inside of your module class, specify all of your ViewModel classes which will participate in this injecting. Instead of having SomeViewModel parameter @Binds to ViewModel, you need to bind the factory to its base class(SomeViewModel.Factory toAssistedSavedStateViewModelFactory<out ViewModel>). You need do this for each ViewModel you want to use the approach with.

@AssistedModule
@Module(includes=[AssistedInject_BuilderModule::class])
abstract class BuilderModule {
  @Binds
  @IntoMap
  @ViewModelKey(SomeViewModel::class)
  abstract fun bindVMFactory(f: SomeViewModel.Factory): AssistedSavedStateViewModelFactory<out ViewModel>
}

5.

As with regular Dagger and ViewModel injecting, you need to have custom factory and binding ViewModel classes to @Multibinds map. For SavedState factory, AndroidX provides necessary implementation — AbstractSavedStateViewModelFactory, which has create() method with the SavedStateHandle parameters to pass into ViewModel. The trick here is not to create the factory directly, but to create an abstract factory (factory creating factories), so that in your Fragment/Activity you can call create() manually and provide the dynamic parameters defaultArguments (Bundle) and this (SavedStateRegistryOwner).

@Reusable
class InjectingSavedStateViewModelFactory @Inject constructor(
    private val assistedFactories: Map<Class<out ViewModel>, @JvmSuppressWildcards AssistedSavedStateViewModelFactory<out ViewModel>>
) {
    fun create(owner: SavedStateRegistryOwner, defaultArgs: Bundle? = null): AbstractSavedStateViewModelFactory {
        return object : AbstractSavedStateViewModelFactory(owner, defaultArgs) {
          
            @Suppress("UNCHECKED_CAST")
            override fun <T : ViewModel?> create(key: String, modelClass: Class<T>, handle: SavedStateHandle): T {
                // Attempt to get ViewModel from assisted inject factories
                assistedFactories[modelClass]?.let {
                    try {
                        return it.create(handle) as T
                    } catch (e: Exception) {
                        throw RuntimeException(e)
                    }
                } ?: throw IllegalArgumentException("Unknown model class $modelClass")
            }
        }
    }
}

6.

In your Fragment, @Inject the InjectingSavedStateViewModelFactory and use it for retrieving your ViewModel class. This way you have all ViewModel's dependencies participating in the Dagger graph, and the possibility to pass dynamic parameters.

class SomeFragment() : Fragment() {
  
  // ...	
	
  @Inject
  lateinit var abstractFactory: InjectingSavedStateViewModelFactory

  lateinit var viewModel: SomeViewModel

  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    // default arguments, so you can set something dynamically
    val defArgs = bundleOf("id" to 5) // may be null
	  
    // now create the actual factory for dagger
    val factory = abstractFactory.create(this, defArgs)
	  
    // get the ViewModel with the factory and scope you want
    viewModel = ViewModelProvider(this, factory)[SomeViewModel::class.java]
  }

  // ...
  
}

You may actually @Inject this factory in your BaseFragment class and save lines of code for each Fragment you inherit.

WANT TO ADD THIS TO AN EXISTING PROJECT?

This step is meant for already running projects, which have Dagger set up with “plain” ViewModels, but where you would like to start using SavedState module. It's also possible that all of your ViewModels won't need to use SavedState in the future; in this case, you can write less boilerplate and have the option of not writing AssistedInject factories.

In usual Dagger + ViewModel setup, you bind your classes to map of ViewModel class and its Provider<ViewModel>. This structure is then injected into your ViewModelFactory, which takes care of instantiating ViewModels.

Map<Class<ViewModel>, @JvmSuppressWildcards Provider<ViewModel>>

In case of AssistedInject, we can see that you need to tweak this structure a bit. You need to have a map of ViewModel class and its AssistedSavedStateViewModelFactory.

Map<Class<out ViewModel>, @JvmSuppressWildcards AssistedSavedStateViewModelFactory<out ViewModel>>

For our multi-purpose factory, we @Inject both structures and let the create method attempt to instantiate ViewModel in each way. And if we're not successful, we use crashing as the last resort.

In the code below, there are two methods to inject — first with @AssistedInject and second with regular @Inject. Replace the previous factory with this multi-purpose one and then you can have both types of ViewModel creators.

@Reusable
class InjectingSavedStateViewModelFactory @Inject constructor(
    private val assistedFactories: Map<Class<out ViewModel>, @JvmSuppressWildcards AssistedSavedStateViewModelFactory<out ViewModel>>,
    private val viewModelProviders: Map<Class<out ViewModel>, @JvmSuppressWildcards Provider<ViewModel>>
) {
    /**
     * Creates instance of ViewModel either annotated with @AssistedInject or @Inject and passes dependencies it needs.
     */
    fun create(owner: SavedStateRegistryOwner, defaultArgs: Bundle? = null) =
        object : AbstractSavedStateViewModelFactory(owner, defaultArgs) {
            override fun <T : ViewModel?> create(
                key: String,
                modelClass: Class<T>,
                handle: SavedStateHandle
            ): T {
                val viewModel =
                    createAssistedInjectViewModel(modelClass, handle)
                        ?: createInjectViewModel(modelClass)
                        ?: throw IllegalArgumentException("Unknown model class $modelClass")

                try {
                    @Suppress("UNCHECKED_CAST")
                    return viewModel as T
                } catch (e: Exception) {
                    throw RuntimeException(e)
                }
            }
        }

    /**
     * Creates ViewModel based on @AssistedInject constructor and its factory
     */
    private fun <T : ViewModel?> createAssistedInjectViewModel(
        modelClass: Class<T>,
        handle: SavedStateHandle
    ): ViewModel? {
        val creator = assistedFactories[modelClass]
            ?: assistedFactories.asIterable().firstOrNull { modelClass.isAssignableFrom(it.key) }?.value
            ?: return null

        return creator.create(handle)
    }

    /**
     * Creates ViewModel based on regular Dagger @Inject constructor
     */
    private fun <T : ViewModel?> createInjectViewModel(modelClass: Class<T>): ViewModel? {
        val creator = viewModelProviders[modelClass]
            ?: viewModelProviders.asIterable().firstOrNull { modelClass.isAssignableFrom(it.key) }?.value
            ?: return null

        return creator.get()
    }
}

TLDR

As promised, for those who wanted to skip the article, here are the steps that make the magic happen:

1. Add gradle dependency.

compileOnly 
  'com.squareup.inject:assisted-inject-annotations-dagger2:0.5.2'
kapt 
  'com.squareup.inject:assisted-inject-processor-dagger2:0.5.2'

2. Update your ViewModel with the AssistedInject library.

class SomeViewModel @AssistedInject constructor(
    private val application: Application,
    @Assisted private val savedStateHandle: SavedStateHandle
){
    // must be inside of the ViewModel class!
    @AssistedInject.Factory
    interface Factory : AssistedSavedStateViewModelFactory<SomeViewModel> {
        override fun create(savedStateHandle: SavedStateHandle): SomeViewModel  // may be ommited prior kotlin 1.3.60 or after PR #121 in AssistedInject lib
    }
}

3. Create (or update your current) Dagger module with @AssistedModule, include generated module AssistedInject_MyDaggerModule::class and @Binds your ViewModel factories @IntoMap.

@AssistedModule
@Module(includes=[AssistedInject_BuilderModule::class])
abstract class BuilderModule {
  @Binds
  @IntoMap
  @ViewModelKey(SomeViewModel::class)
  abstract fun bindVMFactory(f: SomeViewModel.Factory): AssistedSavedStateViewModelFactory<out ViewModel>
}

4. Create (or update your current) ViewModel factory
See InjectingSavedStateViewModelFactory.kt gist.

5. Retrieve the ViewModel with the injecting factory.

CONCLUSION

In this article, I covered how to set up Dagger to work with ViewModels and Saved State module in a fairly boilerplate-less way. Well, the first setup is really boilerplate. But after that, for each ViewModel, it's OK.

In addition to setting up in terms of new projects, there's also the option to update your current Dagger setup and start using it without rewriting the entire existing project.

There's a sample project with one Activity, Fragment and two ViewModels. One with @Inject and one with @AssistedInject.

https://github.com/mlykotom/connecting-the-dots-sample

REFERENCES

The articles below led me to writing this article. Each describes one “dot,” but connecting them all is the hardest part.

https://proandroiddev.com/saving-ui-state-with-viewmodel-savedstate-and-dagger-f77bcaeb8b08

https://proandroiddev.com/brave-new-android-world-with-assistedinject-d11bdc20147d

https://www.coroutinedispatcher.com/2019/08/how-to-produce-savedstatehandle-in-your.html

Thank you to everyone who reviewed this article, especially Marek Abaffy, Michal Urbanek and Iveta Jurčíková.

Share Article
Tomas Mlynaric

Tomas Mlynaric

Android Developer

You might also like...