Add Context.registerClosable (#497)

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 ### Added
 - Added `limit` parameter to `option().counted()` to limit the number of times the option can be used. You can either clamp the value to the limit, or throw an error if the limit is exceeded. ([#483](https://github.com/ajalt/clikt/issues/483))
+- Added `Context.registerClosable` and `Context.callOnClose` to allow you to register cleanup actions that will be called when the command exits. ([#395](https://github.com/ajalt/clikt/issues/395))
 
 ## 4.2.2
 ### Changed

clikt/api/clikt.api

@@ -161,6 +161,8 @@ public final class com/github/ajalt/clikt/core/Context {
 	public static final field Companion Lcom/github/ajalt/clikt/core/Context$Companion;
 	public synthetic fun <init> (Lcom/github/ajalt/clikt/core/Context;Lcom/github/ajalt/clikt/core/CliktCommand;ZZLjava/lang/String;ZLjava/util/Set;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;Lcom/github/ajalt/mordant/terminal/Terminal;Lkotlin/jvm/functions/Function1;ZLcom/github/ajalt/clikt/sources/ValueSource;Lkotlin/jvm/functions/Function2;Lcom/github/ajalt/clikt/output/Localization;Lkotlin/jvm/functions/Function1;Ljava/lang/Object;Ljava/util/List;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun ancestors ()Lkotlin/sequences/Sequence;
+	public final fun callOnClose (Lkotlin/jvm/functions/Function0;)V
+	public final fun close ()V
 	public final fun commandNameWithParents ()Ljava/util/List;
 	public final fun fail (Ljava/lang/String;)Ljava/lang/Void;
 	public static synthetic fun fail$default (Lcom/github/ajalt/clikt/core/Context;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/Void;
@@ -239,8 +241,13 @@ public abstract interface class com/github/ajalt/clikt/core/ContextCliktError {
 	public abstract fun setContext (Lcom/github/ajalt/clikt/core/Context;)V
 }
 
+public final class com/github/ajalt/clikt/core/ContextJvmKt {
+	public static final fun registerJvmCloseable (Lcom/github/ajalt/clikt/core/Context;Ljava/lang/AutoCloseable;)Ljava/lang/AutoCloseable;
+}
+
 public final class com/github/ajalt/clikt/core/ContextKt {
 	public static final fun getTheme (Lcom/github/ajalt/clikt/core/Context;)Lcom/github/ajalt/mordant/rendering/Theme;
+	public static final fun registerCloseable (Lcom/github/ajalt/clikt/core/Context;Ljava/lang/AutoCloseable;)Ljava/lang/AutoCloseable;
 }
 
 public final class com/github/ajalt/clikt/core/FileNotFound : com/github/ajalt/clikt/core/UsageError {

clikt/src/commonMain/kotlin/com/github/ajalt/clikt/core/Context.kt

@@ -129,6 +129,8 @@ class Context private constructor(
     var errorEncountered: Boolean = false
         internal set
 
+    private val closeables = mutableListOf<() -> Unit>()
+
     /** Find the closest object of type [T] */
     inline fun <reified T : Any> findObject(): T? {
         return selfAndAncestors().mapNotNull { it.obj as? T }.firstOrNull()
@@ -163,6 +165,43 @@ class Context private constructor(
     /** Throw a [UsageError] with the given message */
     fun fail(message: String = ""): Nothing = throw UsageError(message)
 
+    /**
+     * Register a callback to be called when this command and all its subcommands have finished.
+     *
+     * This is useful for resources that need to be shared across multiple commands.
+     *
+     * If your resource implements [AutoCloseable], you should use [registerCloseable] instead.
+     *
+     * ### Example
+     *
+     * ```
+     * currentContext.callOnClose { myResource.close() }
+     * ```
+     */
+    fun callOnClose(closeable: () -> Unit) {
+        closeables.add(closeable)
+    }
+
+    /**
+     * Close all registered closeables in the reverse order they were registered.
+     *
+     * This is called automatically after a command and its subcommands have finished running.
+     */
+    fun close() {
+        var err: Throwable? = null
+        for (c in closeables.asReversed()) {
+            try {
+                c()
+            } catch (e: Throwable) {
+                if (err == null) err = e
+                else err.addSuppressed(e)
+            }
+        }
+        closeables.clear()
+        if (err != null) throw err
+    }
+
+    // TODO(5.0): these don't need to be member functions
     @PublishedApi
     internal fun ancestors() = generateSequence(parent) { it.parent }
 
@@ -342,6 +381,30 @@ class Context private constructor(
     }
 }
 
+/**
+ * Register an [AutoCloseable] to be closed when this command and all its subcommands have
+ * finished running.
+ *
+ * This is useful for resources that need to be shared across multiple commands. For resources
+ * that aren't shared, it's often simpler to use [use] directly.
+ *
+ * Registered closeables will be closed in the reverse order that they were registered.
+ *
+ * ### Example
+ *
+ * ```
+ * currentContext.obj = currentContext.registerCloseable(MyResource())
+ * ```
+ *
+ * @return the closeable that was registered
+ * @see Context.callOnClose
+ */
+@ExperimentalStdlibApi
+fun <T: AutoCloseable> Context.registerCloseable(closeable: T): T {
+    callOnClose { closeable.close() }
+    return closeable
+}
+
 /** Find the closest object of type [T], or throw a [NullPointerException] */
 @Suppress("UnusedReceiverParameter") // these extensions don't use their receiver, but we want to limit where they can be called
 inline fun <reified T : Any> CliktCommand.requireObject(): ReadOnlyProperty<CliktCommand, T> {

clikt/src/commonMain/kotlin/com/github/ajalt/clikt/parsers/Parser.kt

@@ -3,6 +3,7 @@ package com.github.ajalt.clikt.parsers
 import com.github.ajalt.clikt.core.*
 import com.github.ajalt.clikt.internal.finalizeParameters
 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.parameters.options.splitOptionPrefix
 
@@ -99,15 +100,13 @@ internal object Parser {
             return context.tokenTransformer(context, token.take(2)) !in optionsByName
         }
 
-        fun addError(e: Err) {
-            errors += e
-            context.errorEncountered = true
-        }
-
         fun consumeParse(tokenIndex: Int, result: OptParseResult) {
             positionalArgs += result.unknown.map { tokenIndex to it }
             invocations += result.known
-            result.err?.let(::addError)
+            result.err?.let {
+                errors += it
+                context.errorEncountered = true
+            }
             i += result.consumed
         }
 
@@ -211,92 +210,143 @@ internal object Parser {
         }
 
 
-        // Finalize and validate everything as long as we aren't resuming a parse for multiple subcommands
+        // Finalize and validate everything as long as we aren't resuming a parse for multiple
+        // subcommands
         try {
-            if (canRun) {
-                // Finalize and validate eager options
-                invocationsByOption.forEach { (o, inv) ->
-                    if (o.eager) {
-                        o.finalize(context, inv)
-                        o.postValidate(context)
-                    }
+            try {
+                if (canRun) {
+                    i = finalizeAndRun(
+                        context,
+                        i,
+                        command,
+                        subcommand,
+                        invocationsByOption,
+                        positionalArgs,
+                        arguments,
+                        hasMultipleSubAncestor,
+                        tokens,
+                        subcommands,
+                        errors,
+                        ungroupedOptions,
+                        invocationsByOptionByGroup
+                    )
+                } else if (subcommand == null && positionalArgs.isNotEmpty()) {
+                    // If we're resuming a parse with multiple subcommands, there can't be any args
+                    // after the last subcommand is parsed
+                    throw excessArgsError(positionalArgs, positionalArgs.size, context)
                 }
+            } catch (e: UsageError) {
+                // Augment usage errors with the current context if they don't have one
+                e.context = context
+                throw e
+            }
 
-                // Parse arguments
-                val argsParseResult = parseArguments(i, positionalArgs, arguments)
-                argsParseResult.err?.let(::addError)
+            if (subcommand != null) {
+                val nextTokens = parse(tokens.drop(i), subcommand.currentContext, true)
+                if (command.allowMultipleSubcommands && nextTokens.isNotEmpty()) {
+                    parse(nextTokens, context, false)
+                }
+                return nextTokens
+            }
+        } finally {
+            context.close()
+        }
 
-                val excessResult = handleExcessArguments(
-                    argsParseResult.excessCount,
-                    hasMultipleSubAncestor,
-                    i,
-                    tokens,
-                    subcommands,
-                    positionalArgs,
-                    context
-                )
-                excessResult.second?.let(::addError)
+        return tokens.drop(i)
+    }
 
-                val usageErrors = errors
-                    .filter { it.includeInMulti }.ifEmpty { errors }
-                    .sortedBy { it.i }.mapTo(mutableListOf()) { it.e }
+    private fun finalizeAndRun(
+        context: Context,
+        i: Int,
+        command: CliktCommand,
+        subcommand: CliktCommand?,
+        invocationsByOption: Map<Option, List<Invocation>>,
+        positionalArgs: MutableList<Pair<Int, String>>,
+        arguments: MutableList<Argument>,
+        hasMultipleSubAncestor: Boolean,
+        tokens: List<String>,
+        subcommands: Map<String, CliktCommand>,
+        errors: MutableList<Err>,
+        ungroupedOptions: List<Option>,
+        invocationsByOptionByGroup: Map<ParameterGroup?, Map<Option, List<Invocation>>>,
+    ): Int {
+        // Finalize and validate eager options
+        var nextArgvI = i
+
+        invocationsByOption.forEach { (o, inv) ->
+            if (o.eager) {
+                o.finalize(context, inv)
+                o.postValidate(context)
+            }
+        }
 
-                i = excessResult.first
+        // Parse arguments
+        val argsParseResult = parseArguments(nextArgvI, positionalArgs, arguments)
+        argsParseResult.err?.let {
+            errors += it
+            context.errorEncountered = true
+        }
 
-                // Finalize arguments, groups, and options
-                gatherErrors(usageErrors, context) {
-                    finalizeParameters(
-                        context,
-                        ungroupedOptions,
-                        command._groups,
-                        invocationsByOptionByGroup,
-                        argsParseResult.args,
-                    )
-                }
+        val excessResult = handleExcessArguments(
+            argsParseResult.excessCount,
+            hasMultipleSubAncestor,
+            nextArgvI,
+            tokens,
+            subcommands,
+            positionalArgs,
+            context
+        )
+        excessResult.second?.let {
+            errors += it
+            context.errorEncountered = true
+        }
 
-                // 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 }
+        val usageErrors = errors
+            .filter { it.includeInMulti }.ifEmpty { errors }
+            .sortedBy { it.i }.mapTo(mutableListOf()) { it.e }
 
-                // Now that all parameters have been finalized, we can validate everything
-                ungroupedOptions.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
-                command._groups.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
-                command._arguments.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
+        nextArgvI = excessResult.first
 
-                MultiUsageError.buildOrNull(usageErrors)?.let { throw it }
+        // Finalize arguments, groups, and options
+        gatherErrors(usageErrors, context) {
+            finalizeParameters(
+                context,
+                ungroupedOptions,
+                command._groups,
+                invocationsByOptionByGroup,
+                argsParseResult.args,
+            )
+        }
 
-                if (subcommand == null && subcommands.isNotEmpty() && !command.invokeWithoutSubcommand) {
-                    throw PrintHelpMessage(context, error = true)
-                }
+        // 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 }
+
+        // Now that all parameters have been finalized, we can validate everything
+        ungroupedOptions.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
+        command._groups.forEach { gatherErrors(usageErrors, context) { it.postValidate(context) } }
+        command._arguments.forEach {
+            gatherErrors(
+                usageErrors,
+                context
+            ) { it.postValidate(context) }
+        }
 
-                command.currentContext.invokedSubcommand = subcommand
-                if (command.currentContext.printExtraMessages) {
-                    for (warning in command.messages) {
-                        command.terminal.warning(warning, stderr = true)
-                    }
-                }
+        MultiUsageError.buildOrNull(usageErrors)?.let { throw it }
 
-                command.run()
-            } else if (subcommand == null && positionalArgs.isNotEmpty()) {
-                // If we're resuming a parse with multiple subcommands, there can't be any args after the last
-                // subcommand is parsed
-                throw excessArgsError(positionalArgs, positionalArgs.size, context)
-            }
-        } catch (e: UsageError) {
-            // Augment usage errors with the current context if they don't have one
-            e.context = context
-            throw e
+        if (subcommand == null && subcommands.isNotEmpty() && !command.invokeWithoutSubcommand) {
+            throw PrintHelpMessage(context, error = true)
         }
 
-        if (subcommand != null) {
-            val nextTokens = parse(tokens.drop(i), subcommand.currentContext, true)
-            if (command.allowMultipleSubcommands && nextTokens.isNotEmpty()) {
-                parse(nextTokens, context, false)
+        command.currentContext.invokedSubcommand = subcommand
+        if (command.currentContext.printExtraMessages) {
+            for (warning in command.messages) {
+                command.terminal.warning(warning, stderr = true)
             }
-            return nextTokens
         }
 
-        return tokens.drop(i)
+        command.run()
+        return nextArgvI
     }
 
     /** Returns either the new argv index, or an error */
@@ -514,14 +564,14 @@ internal object Parser {
             1 -> context.localization.extraArgumentOne(actual)
             else -> context.localization.extraArgumentMany(actual, excess)
         }
-        return UsageError(message)
+        return UsageError(message).also { it.context = context }
     }
 }
 
 private inline fun gatherErrors(
     errors: MutableList<UsageError>,
     context: Context,
-    block: () -> Unit
+    block: () -> Unit,
 ) {
     try {
         block()