From b3a145ab31d1b9c57c26745685c17e4198978033 Mon Sep 17 00:00:00 2001
From: AJ <aj@alt>
Date: Sun, 17 Dec 2023 09:49:34 -0800
Subject: [PATCH] Improve support for recursive lazy paraemeters

---
 .../ajalt/clikt/internal/Finalization.kt      | 91 ++++++++++++-------
 .../com/github/ajalt/clikt/parsers/Parser.kt  | 72 +++------------
 .../ajalt/clikt/core/CliktCommandTest.kt      |  1 +
 .../ajalt/clikt/parameters/ArgumentTest.kt    |  8 +-
 .../ajalt/clikt/parameters/OptionTest.kt      | 13 +++
 .../parameters/groups/OptionGroupsTest.kt     | 21 ++++-
 docs/advanced.md                              |  8 +-
 docs/options.md                               | 14 ++-
 8 files changed, 118 insertions(+), 110 deletions(-)

diff --git a/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/internal/Finalization.kt b/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/internal/Finalization.kt
index db896a1e1..e165aa9ea 100644
--- a/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/internal/Finalization.kt
+++ b/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/internal/Finalization.kt
@@ -4,6 +4,8 @@ import com.github.ajalt.clikt.core.Abort
 import com.github.ajalt.clikt.core.Context
 import com.github.ajalt.clikt.core.MultiUsageError
 import com.github.ajalt.clikt.core.UsageError
+import com.github.ajalt.clikt.parameters.arguments.Argument
+import com.github.ajalt.clikt.parameters.groups.ParameterGroup
 import com.github.ajalt.clikt.parameters.options.Option
 import com.github.ajalt.clikt.parsers.Invocation
 
@@ -12,45 +14,66 @@ internal fun finalizeOptions(
     options: List<Option>,
     invocationsByOption: Map<Option, List<Invocation>>,
 ) {
-    val errors = mutableListOf<UsageError>()
-    // Finalize invoked options
-    for ((option, invocations) in invocationsByOption) {
-        try {
-            option.finalize(context, invocations)
-        } catch (e: UsageError) {
-            errors += e
-            context.errorEncountered = true
-        }
+    finalizeParameters(
+        context, options, emptyList(), mapOf(null to invocationsByOption), emptyList()
+    )
+}
+
+private sealed class Param
+private class Opt(val option: Option, val invs: List<Invocation>) : Param()
+private class Arg(val argument: Argument, val invs: List<String>) : Param()
+private class Group(val group: ParameterGroup, val invs: Map<Option, List<Invocation>>) : Param()
+
+internal fun finalizeParameters(
+    context: Context,
+    options: List<Option>,
+    groups: List<ParameterGroup>,
+    invocationsByOptionByGroup: Map<ParameterGroup?, Map<Option, List<Invocation>>>,
+    argsAndTokens: List<Pair<Argument, List<String>>>,
+) {
+    // Add uninvoked params last so that e.g. we can skip prompting if there's an error in an
+    // invoked option
+    val allGroups = invocationsByOptionByGroup.toMutableMap()
+    for (group in groups) {
+        if (group !in allGroups) allGroups[group] = emptyMap()
     }
 
-    // Finalize uninvoked options
-    val retries = mutableListOf<Option>()
-    for (o in options.filter { it !in invocationsByOption }) {
-        try {
-            o.finalize(context, emptyList())
-        } catch (e: IllegalStateException) {
-            retries += o
-        } catch (e: UsageError) {
-            errors += e
-        } catch (e: Abort) {
-            // ignore Abort if we already encountered an error
-            if (!context.errorEncountered) throw e
-        }
+    val allOptions = (invocationsByOptionByGroup[null] ?: emptyMap()).toMutableMap()
+    for (opt in options) {
+        if (opt !in allOptions) allOptions[opt] = emptyList()
     }
 
-    // If an uninvoked option triggers an ISE, retry it once after other options have been finalized
-    // so that lazy defaults can reference other options.
-    for (o in retries) {
-        try {
-            o.finalize(context, emptyList())
-        } catch (e: IllegalStateException) {
-            // If we had an earlier usage error, throw that instead of the ISE
-            MultiUsageError.buildOrNull(errors)?.let { throw it }
-            throw e
-        } catch (e: UsageError) {
-            errors += e
-            context.errorEncountered = true
+    val allParams: List<Param> =
+        argsAndTokens.map { Arg(it.first, it.second) } +
+                allOptions.map { Opt(it.key, it.value) } +
+                allGroups.mapNotNull { it.key?.let { k -> Group(k, it.value) } }
+
+
+    val errors = mutableListOf<UsageError>()
+    var currentRound = allParams.toList()
+    val nextRound = mutableListOf<Param>()
+
+    while (true) {
+        for (it in currentRound) {
+            try {
+                when (it) {
+                    is Arg -> it.argument.finalize(context, it.invs)
+                    is Opt -> it.option.finalize(context, it.invs)
+                    is Group -> it.group.finalize(context, it.invs)
+                }
+            } catch (e: IllegalStateException) {
+                nextRound += it
+            } catch (e: UsageError) {
+                errors += e
+                context.errorEncountered = true
+            } catch (e: Abort) {
+                // ignore Abort if we already encountered an error
+                if (!context.errorEncountered) throw e
+            }
         }
+        if (currentRound.size <= nextRound.size) break
+        currentRound = nextRound.toList()
+        nextRound.clear()
     }
 
     MultiUsageError.buildOrNull(errors)?.let { throw it }
diff --git a/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/parsers/Parser.kt b/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/parsers/Parser.kt
index ba16c0321..a73aa8891 100644
--- a/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/parsers/Parser.kt
+++ b/clikt/src/commonMain/kotlin/com/github/ajalt/clikt/parsers/Parser.kt
@@ -1,7 +1,7 @@
 package com.github.ajalt.clikt.parsers
 
 import com.github.ajalt.clikt.core.*
-import com.github.ajalt.clikt.internal.finalizeOptions
+import com.github.ajalt.clikt.internal.finalizeParameters
 import com.github.ajalt.clikt.parameters.arguments.Argument
 import com.github.ajalt.clikt.parameters.options.Option
 import com.github.ajalt.clikt.parameters.options.splitOptionPrefix
@@ -11,7 +11,8 @@ private data class Err(val e: UsageError, val i: Int, val includeInMulti: Boolea
 
 private data class ArgsParseResult(
     val excessCount: Int,
-    val args: Map<Argument, List<String>>,
+    /** All arguments in the command and their parsed tokens (if any) */
+    val args: List<Pair<Argument, List<String>>>,
     val err: Err?,
 )
 
@@ -242,48 +243,20 @@ internal object Parser {
 
                 i = excessResult.first
 
-                // Finalize arguments before options, so that options can reference them
-                val (retries, argErrs) = finalizeArguments(argsParseResult.args, context)
-                usageErrors += argErrs
-                if (argErrs.isNotEmpty()) context.errorEncountered = true
-
-                // Finalize un-grouped options
+                // Finalize arguments, groups, and options
                 gatherErrors(usageErrors, context) {
-                    finalizeOptions(
+                    finalizeParameters(
                         context,
                         ungroupedOptions,
-                        invocationsByOptionByGroup[null] ?: emptyMap()
+                        command._groups,
+                        invocationsByOptionByGroup,
+                        argsParseResult.args,
                     )
                 }
 
-                // Since groups like ChoiceGroup may depend on an ungrouped option's value, we can't finalize groups if
-                // any options failed to validate, so throw errors now
-                MultiUsageError.buildOrNull(usageErrors)?.let { throw it }
-                usageErrors.clear()
-
-                // Finalize groups after ungrouped options so that the groups can use their values
-                for ((group, invs) in invocationsByOptionByGroup) {
-                    gatherErrors(usageErrors, context) { group?.finalize(context, invs) }
-                }
-
-                // Finalize groups with no invocations
-                for (g in command._groups) {
-                    if (g !in invocationsByGroup) gatherErrors(usageErrors, context) {
-                        g.finalize(context, emptyMap())
-                    }
-                }
-
                 // We can't validate a param that didn't finalize successfully, and we don't keep
                 // track of which ones are finalized, so throw any errors now
                 MultiUsageError.buildOrNull(usageErrors)?.let { throw it }
-                usageErrors.clear()
-
-                // Retry any failed args now that they can reference option values
-                retries.forEach { (arg, v) ->
-                    gatherErrors(usageErrors, context) {
-                        arg.finalize(context, v)
-                    }
-                }
 
                 // Now that all parameters have been finalized, we can validate everything
                 ungroupedOptions.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
@@ -485,6 +458,7 @@ internal object Parser {
 
     private fun parseArguments(
         argvIndex: Int,
+        /** Parsed argument tokens and their index in argv */
         positionalArgs: List<Pair<Int, String>>,
         arguments: List<Argument>,
     ): ArgsParseResult {
@@ -512,7 +486,7 @@ internal object Parser {
                 }
                 return ArgsParseResult(
                     0,
-                    out,
+                    out.toList(),
                     Err(e, positionalArgs.getOrNull(i)?.first ?: argvIndex)
                 )
             }
@@ -522,33 +496,9 @@ internal object Parser {
         }
 
         val excess = positionalArgs.size - i
-        return ArgsParseResult(excess, out, null)
-    }
-
-    /** Returns map of arguments that need retries to their values */
-    private fun finalizeArguments(
-        parsedArgs: Map<Argument, List<String>>,
-        context: Context,
-    ): Pair<Map<Argument, List<String>>, List<UsageError>> {
-        val retries = mutableMapOf<Argument, List<String>>()
-        val usageErrors = mutableListOf<UsageError>()
-        for ((it, v) in parsedArgs) {
-            try {
-                it.finalize(context, v)
-            } catch (e: IllegalStateException) {
-                retries[it] = v
-            } catch (e: UsageError) {
-                usageErrors += e
-                context.errorEncountered = true
-            } catch (e: Abort) {
-                // ignore Abort if we already encountered an error
-                if (!context.errorEncountered) throw e
-            }
-        }
-        return retries to usageErrors
+        return ArgsParseResult(excess, out.toList(), null)
     }
 
-
     private fun loadArgFile(filename: String, context: Context): List<String> {
         return shlex(filename, context.argumentFileReader!!(filename), context)
     }
diff --git a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/core/CliktCommandTest.kt b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/core/CliktCommandTest.kt
index 0c18c46e4..bbf4a7e77 100644
--- a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/core/CliktCommandTest.kt
+++ b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/core/CliktCommandTest.kt
@@ -6,6 +6,7 @@ import com.github.ajalt.clikt.parameters.groups.OptionGroup
 import com.github.ajalt.clikt.parameters.groups.cooccurring
 import com.github.ajalt.clikt.parameters.groups.mutuallyExclusiveOptions
 import com.github.ajalt.clikt.parameters.groups.provideDelegate
+import com.github.ajalt.clikt.parameters.options.defaultLazy
 import com.github.ajalt.clikt.parameters.options.flag
 import com.github.ajalt.clikt.parameters.options.option
 import com.github.ajalt.clikt.parameters.options.required
diff --git a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/ArgumentTest.kt b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/ArgumentTest.kt
index d5d645ba2..6dda101c5 100644
--- a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/ArgumentTest.kt
+++ b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/ArgumentTest.kt
@@ -51,17 +51,15 @@ class ArgumentTest {
 
     @Test
     @JsName("one_default_argument")
-    fun `one default argument`() = forAll(
-        row("", "def")
-    ) { argv, expected ->
+    fun `one default argument`() {
         class C : TestCommand() {
             val x by argument().default("def")
             override fun run_() {
-                x shouldBe expected
+                x shouldBe "def"
             }
         }
 
-        C().parse(argv)
+        C().parse("")
     }
 
     @Test
diff --git a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/OptionTest.kt b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/OptionTest.kt
index eeb03da4d..b07e7c454 100644
--- a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/OptionTest.kt
+++ b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/OptionTest.kt
@@ -455,6 +455,19 @@ class OptionTest {
         C().parse("")
     }
 
+    @Test
+    @JsName("defaultLazy_option_mutual_recursion")
+    fun `defaultLazy option mutual recursion`() {
+        class C : TestCommand() {
+            val y: String by option().defaultLazy { x }
+            val x: String by option().defaultLazy { y }
+            override fun run_() {
+                y shouldBe "def"
+            }
+        }
+        shouldThrow<IllegalStateException> { C().parse("") }
+    }
+
     @Test
     @JsName("required_option")
     fun `required option`() {
diff --git a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/groups/OptionGroupsTest.kt b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/groups/OptionGroupsTest.kt
index 6841ed184..96a30f878 100644
--- a/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/groups/OptionGroupsTest.kt
+++ b/clikt/src/commonTest/kotlin/com/github/ajalt/clikt/parameters/groups/OptionGroupsTest.kt
@@ -2,15 +2,13 @@
 
 package com.github.ajalt.clikt.parameters.groups
 
-import com.github.ajalt.clikt.core.BadParameterValue
-import com.github.ajalt.clikt.core.MissingOption
-import com.github.ajalt.clikt.core.MutuallyExclusiveGroupException
-import com.github.ajalt.clikt.core.UsageError
+import com.github.ajalt.clikt.core.*
 import com.github.ajalt.clikt.parameters.options.*
 import com.github.ajalt.clikt.parameters.types.int
 import com.github.ajalt.clikt.testing.TestCommand
 import com.github.ajalt.clikt.testing.formattedMessage
 import com.github.ajalt.clikt.testing.parse
+import com.github.ajalt.clikt.testing.test
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.data.blocking.forAll
 import io.kotest.data.row
@@ -579,6 +577,21 @@ class OptionGroupsTest {
         shouldThrow<UsageError> { C(3).parse("--g=a --opt=3") }
         shouldThrow<UsageError> { C(3).parse("--opt=3") }
     }
+
+    @Test
+    @JsName("defaultLazy_option_referencing_group")
+    fun `defaultLazy option referencing group`() {
+        class C : TestCommand() {
+            val g by Group1()
+            val o by option().int().defaultLazy { g.g11 }
+
+            override fun run_() {
+                o shouldBe 1
+            }
+        }
+
+        C().parse("--g11=1")
+    }
 }
 
 private class Group1 : OptionGroup() {
diff --git a/docs/advanced.md b/docs/advanced.md
index 85a9241a7..d1e98fac2 100644
--- a/docs/advanced.md
+++ b/docs/advanced.md
@@ -186,9 +186,11 @@ To implement git-style aliases:
     Committing with message: my message
     ```
 
-Note that aliases are not expanded recursively: none of the tokens that
-an alias expands to will be expanded again, even if they match another
-alias.
+!!! note
+
+    Aliases are not expanded recursively: none of the tokens that
+    an alias expands to will be expanded again, even if they match another
+    alias.
 
 You also use this functionality to implement command prefixes:
 
diff --git a/docs/options.md b/docs/options.md
index 969e0bde6..b8d83bae2 100644
--- a/docs/options.md
+++ b/docs/options.md
@@ -140,6 +140,12 @@ you can use [`defaultLazy`][defaultLazy] instead of [`default`][default]. It has
 but you give it a lambda returning the default value, and the lambda will only be called if the
 option isn't present on the command line.
 
+!!! warning
+
+    The lambda you pass to `defaultLazy` is normally called at most once. But the lambda references
+    another parameter, it may be called more than once in order to make sure the parameter it references
+    is available.
+
 ## Multi Value Options
 
 Options can take a fixed number of values separated by whitespace, a variable number of values separated by a delimiter,
@@ -1178,9 +1184,11 @@ options:
     Hello, Foo!
     ```
 
-Note that inferred names will always have a POSIX-style prefix like
-`--name`. If you want to use a different prefix, you should specify all
-option names manually.
+!!! note
+
+    Inferred names will always have a POSIX-style prefix like
+    `--name`. If you want to use a different prefix, you should specify all
+    option names manually.
 
 ## Option Transformation Order