From 011081488533b527279e6c14b51a92d9a5e210d3 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Tue, 3 Jan 2023 18:38:17 +0900
Subject: [PATCH] cli: integrate custom StoreFactories registration with
 CliRunner

.set_store_factories() could be invoked for any CliRunner<F> type, but
I don't care much about that since the type parameter F will be removed
soon.
---
 examples/custom-backend/main.rs | 16 ++++++++++------
 src/cli_util.rs                 | 19 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/examples/custom-backend/main.rs b/examples/custom-backend/main.rs
index 726807b84..812103e77 100644
--- a/examples/custom-backend/main.rs
+++ b/examples/custom-backend/main.rs
@@ -34,11 +34,7 @@ enum CustomCommands {
     InitJit,
 }
 
-fn run(
-    ui: &mut Ui,
-    mut command_helper: CommandHelper,
-    matches: &ArgMatches,
-) -> Result<(), CommandError> {
+fn create_store_factories() -> StoreFactories {
     let mut store_factories = StoreFactories::default();
     // Register the backend so it can be loaded when the repo is loaded. The name
     // must match `Backend::name()`.
@@ -46,7 +42,14 @@ fn run(
         "jit",
         Box::new(|store_path| Box::new(JitBackend::load(store_path))),
     );
-    command_helper.set_store_factories(store_factories);
+    store_factories
+}
+
+fn run(
+    ui: &mut Ui,
+    command_helper: CommandHelper,
+    matches: &ArgMatches,
+) -> Result<(), CommandError> {
     match CustomCommands::from_arg_matches(matches) {
         // Handle our custom command
         Ok(CustomCommands::InitJit) => {
@@ -64,6 +67,7 @@ fn run(
 
 fn main() {
     CliRunner::init()
+        .set_store_factories(create_store_factories())
         .add_subcommand::<CustomCommands>()
         .set_dispatch_fn(run)
         .run_and_exit();
diff --git a/src/cli_util.rs b/src/cli_util.rs
index c4a4da222..31c658a20 100644
--- a/src/cli_util.rs
+++ b/src/cli_util.rs
@@ -1688,6 +1688,7 @@ pub fn handle_command_result(ui: &mut Ui, result: Result<(), CommandError>) -> i
 pub struct CliRunner<F> {
     tracing_subscription: TracingSubscription,
     app: clap::Command,
+    store_factories: Option<StoreFactories>,
     dispatch_fn: F,
 }
 
@@ -1700,10 +1701,21 @@ impl CliRunner<()> {
         CliRunner {
             tracing_subscription,
             app: crate::commands::default_app(),
+            store_factories: None,
             dispatch_fn: (),
         }
     }
 
+    /// Replaces `StoreFactories` to be used.
+    pub fn set_store_factories(self, store_factories: StoreFactories) -> Self {
+        CliRunner {
+            tracing_subscription: self.tracing_subscription,
+            app: self.app,
+            store_factories: Some(store_factories),
+            dispatch_fn: self.dispatch_fn,
+        }
+    }
+
     /// Registers new subcommands in addition to the default ones.
     // TODO: maybe take dispatch_fn for the subcommands?
     pub fn add_subcommand<C>(self) -> Self
@@ -1713,6 +1725,7 @@ impl CliRunner<()> {
         CliRunner {
             tracing_subscription: self.tracing_subscription,
             app: C::augment_subcommands(self.app),
+            store_factories: self.store_factories,
             dispatch_fn: self.dispatch_fn,
         }
     }
@@ -1725,6 +1738,7 @@ impl CliRunner<()> {
         CliRunner {
             tracing_subscription: self.tracing_subscription,
             app: self.app,
+            store_factories: self.store_factories,
             dispatch_fn,
         }
     }
@@ -1736,12 +1750,15 @@ where
 {
     pub fn run(self, ui: &mut Ui) -> Result<(), CommandError> {
         ui.reset(crate::config::read_config()?);
-        let (command_helper, matches) = parse_args(
+        let (mut command_helper, matches) = parse_args(
             ui,
             self.app,
             &self.tracing_subscription,
             std::env::args_os(),
         )?;
+        if let Some(store_factories) = self.store_factories {
+            command_helper.set_store_factories(store_factories);
+        }
         // TODO: pass CommandHelper by reference
         (self.dispatch_fn)(ui, command_helper, &matches)
     }