provide more flexible mechanism for directive handlers: possibility to pass next handler and invoke it or not

Author: adamsitnik

src/System.CommandLine.ApiCompatibility.Tests/ApiCompatibilityApprovalTests.System_CommandLine_api_is_not_changed.approved.txt

@@ -108,10 +108,10 @@ System.CommandLine
     public static System.Void Write(this IConsole console, System.String value)
     public static System.Void WriteLine(this IConsole console, System.String value)
   public class Directive : Symbol
-    .ctor(System.String name, System.String description = null, System.Action<System.CommandLine.Invocation.InvocationContext> syncHandler = null, System.Func<System.CommandLine.Invocation.InvocationContext,System.Threading.CancellationToken,System.Threading.Tasks.Task> asyncHandler = null)
+    .ctor(System.String name, System.String description = null, System.Action<System.CommandLine.Invocation.InvocationContext,ICommandHandler> syncHandler = null, System.Func<System.CommandLine.Invocation.InvocationContext,ICommandHandler,System.Threading.CancellationToken,System.Threading.Tasks.Task> asyncHandler = null)
     public System.Collections.Generic.IEnumerable<System.CommandLine.Completions.CompletionItem> GetCompletions(System.CommandLine.Completions.CompletionContext context)
-    public System.Void SetAsynchronousHandler(System.Func<System.CommandLine.Invocation.InvocationContext,System.Threading.CancellationToken,System.Threading.Tasks.Task> handler)
-    public System.Void SetSynchronousHandler(System.Action<System.CommandLine.Invocation.InvocationContext> handler)
+    public System.Void SetAsynchronousHandler(System.Func<System.CommandLine.Invocation.InvocationContext,ICommandHandler,System.Threading.CancellationToken,System.Threading.Tasks.Task> handler)
+    public System.Void SetSynchronousHandler(System.Action<System.CommandLine.Invocation.InvocationContext,ICommandHandler> handler)
   public class EnvironmentVariablesDirective : Directive
     .ctor()
   public static class Handler

src/System.CommandLine.Tests/DirectiveTests.cs

@@ -2,7 +2,9 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Collections.Generic;
+using System.Globalization;
 using System.Linq;
+using System.Threading.Tasks;
 using FluentAssertions;
 using Xunit;
 
@@ -137,6 +139,50 @@ public void When_directives_are_not_enabled_they_are_treated_as_regular_tokens()
                   .BeEquivalentTo("[hello]");
         }
 
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public async Task Directive_can_restore_the_state_after_running_continuation(bool async)
+        {
+            const string plCulture = "pl-PL", enUsCulture = "en-US";
+            const string envVarName = "uniqueName", envVarValue = "just";
+
+            var before = CultureInfo.CurrentUICulture;
+
+            try
+            {
+                CultureInfo.CurrentUICulture = new(enUsCulture);
+
+                bool invoked = false;
+                Option<bool> option = new("-a");
+                RootCommand root = new() { option };
+                CommandLineBuilder builder = new(root);
+                builder.Directives.Add(new EnvironmentVariablesDirective());
+                builder.Directives.Add(new CultureDirective());
+                root.SetHandler(ctx =>
+                {
+                    invoked = true;
+                    CultureInfo.CurrentUICulture.Name.Should().Be(plCulture);
+                    Environment.GetEnvironmentVariable(envVarName).Should().Be(envVarValue);
+                });
+
+                if (async)
+                {
+                    await builder.Build().InvokeAsync($"[culture:{plCulture}] [env:{envVarName}={envVarValue}]");
+                }
+                else
+                {
+                    builder.Build().Invoke($"[culture:{plCulture}] [env:{envVarName}={envVarValue}]");
+                }
+
+                invoked.Should().BeTrue();
+            }
+            finally
+            {
+                CultureInfo.CurrentUICulture = before;
+            }
+        }
+
         private static ParseResult Parse(Option option, Directive directive, string commandLine)
         {
             RootCommand root = new() { option };
@@ -145,5 +191,47 @@ private static ParseResult Parse(Option option, Directive directive, string comm
 
             return root.Parse(commandLine, builder.Build());
         }
+
+        private sealed class CultureDirective : Directive
+        {
+            public CultureDirective() : base("culture")
+            {
+                SetSynchronousHandler((ctx, next) =>
+                {
+                    CultureInfo cultureBefore = CultureInfo.CurrentUICulture;
+
+                    try
+                    {
+                        string cultureName = ctx.ParseResult.FindResultFor(this).Values.Single();
+
+                        CultureInfo.CurrentUICulture = new CultureInfo(cultureName);
+
+                        next?.Invoke(ctx);
+                    }
+                    finally
+                    {
+                        CultureInfo.CurrentUICulture = cultureBefore;
+                    }
+                });
+                SetAsynchronousHandler(async (ctx, next, ct) =>
+                {
+                    CultureInfo cultureBefore = CultureInfo.CurrentUICulture;
+
+                    try
+                    {
+                        string cultureName = ctx.ParseResult.FindResultFor(this).Values.Single();
+
+                        CultureInfo.CurrentUICulture = new CultureInfo(cultureName);
+
+                        await next?.InvokeAsync(ctx, ct);
+                    }
+                    finally
+                    {
+                        CultureInfo.CurrentUICulture = cultureBefore;
+                    }
+                });
+            }
+        }
+
     }
 }

src/System.CommandLine/Directive.cs

@@ -3,7 +3,6 @@
 using System.CommandLine.Invocation;
 using System.Threading.Tasks;
 using System.Threading;
-using System.CommandLine.Parsing;
 
 namespace System.CommandLine
 {
@@ -18,17 +17,22 @@ namespace System.CommandLine
     /// </summary>
     public class Directive : Symbol
     {
+        internal Action<InvocationContext, ICommandHandler?>? SyncHandler;
+        internal Func<InvocationContext, ICommandHandler?, CancellationToken, Task>? AsyncHandler;
+
         /// <summary>
         /// Initializes a new instance of the Directive class.
         /// </summary>
         /// <param name="name">The name of the directive. It can't contain whitespaces.</param>
         /// <param name="description">The description of the directive, shown in help.</param>
         /// <param name="syncHandler">The synchronous action that is invoked when directive is parsed.</param>
         /// <param name="asyncHandler">The asynchronous action that is invoked when directive is parsed.</param>
+        /// <remarks>The second argument of both handlers is next handler than can be invoked.
+        /// Example: a custom directive might just change current culture and run actual command afterwards.</remarks>
         public Directive(string name, 
             string? description = null, 
-            Action<InvocationContext>? syncHandler = null,
-            Func<InvocationContext, CancellationToken, Task>? asyncHandler = null)
+            Action<InvocationContext, ICommandHandler?>? syncHandler = null,
+            Func<InvocationContext, ICommandHandler?, CancellationToken, Task>? asyncHandler = null)
         {
             if (string.IsNullOrWhiteSpace(name))
             {
@@ -46,37 +50,17 @@ public Directive(string name,
             Name = name;
             Description = description;
 
-            if (syncHandler is not null)
-            {
-                SetSynchronousHandler(syncHandler);
-            }
-            else if (asyncHandler is not null)
-            {
-                SetAsynchronousHandler(asyncHandler);
-            }
+            SyncHandler = syncHandler;
+            AsyncHandler = asyncHandler;
         }
 
-        public void SetAsynchronousHandler(Func<InvocationContext, CancellationToken, Task> handler)
-        {
-            if (handler is null)
-            {
-                throw new ArgumentNullException(nameof(handler));
-            }
+        internal bool HasHandler => SyncHandler != null || AsyncHandler != null;
 
-            Handler = new AnonymousCommandHandler(handler);
-        }
-
-        public void SetSynchronousHandler(Action<InvocationContext> handler)
-        {
-            if (handler is null)
-            {
-                throw new ArgumentNullException(nameof(handler));
-            }
-
-            Handler = new AnonymousCommandHandler(handler);
-        }
+        public void SetAsynchronousHandler(Func<InvocationContext, ICommandHandler?, CancellationToken, Task> handler)
+            => AsyncHandler = handler ?? throw new ArgumentNullException(nameof(handler));
 
-        internal ICommandHandler? Handler { get; private set; }
+        public void SetSynchronousHandler(Action<InvocationContext, ICommandHandler?> handler)
+            => SyncHandler = handler ?? throw new ArgumentNullException(nameof(handler));
 
         private protected override string DefaultName => throw new NotImplementedException();
 

src/System.CommandLine/EnvironmentVariablesDirective.cs

@@ -1,5 +1,5 @@
-using System.CommandLine.Invocation;
-using System.CommandLine.Parsing;
+using System.CommandLine.Parsing;
+using System.Threading.Tasks;
 
 namespace System.CommandLine
 {
@@ -10,12 +10,28 @@ public sealed class EnvironmentVariablesDirective : Directive
     {
         public EnvironmentVariablesDirective() : base("env")
         {
-            SetSynchronousHandler(SyncHandler);
+            SetSynchronousHandler((context, next) =>
+            {
+                SetEnvVars(context.ParseResult);
+
+                next?.Invoke(context);
+            });
+            SetAsynchronousHandler((context, next, cancellationToken) =>
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return Task.FromCanceled(cancellationToken);
+                }
+
+                SetEnvVars(context.ParseResult);
+
+                return next?.InvokeAsync(context, cancellationToken) ?? Task.CompletedTask;
+            });
         }
 
-        private void SyncHandler(InvocationContext context)
+        private void SetEnvVars(ParseResult parseResult)
         {
-            DirectiveResult directiveResult = context.ParseResult.FindResultFor(this)!;
+            DirectiveResult directiveResult = parseResult.FindResultFor(this)!;
 
             for (int i = 0; i < directiveResult.Values.Count; i++)
             {
@@ -35,9 +51,6 @@ private void SyncHandler(InvocationContext context)
                     }
                 }
             }
-
-            // we need a cleaner, more flexible and intuitive way of continuing the execution
-            context.ParseResult.CommandResult.Command.Handler?.Invoke(context);
         }
     }
 }

src/System.CommandLine/Invocation/CombinedCommandHandler.cs

@@ -0,0 +1,81 @@
+using System.CommandLine.Parsing;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.CommandLine.Invocation
+{
+    internal sealed class ChainedCommandHandler : ICommandHandler
+    {
+        private readonly SymbolResultTree _symbols;
+        private readonly ICommandHandler? _commandHandler;
+
+        internal ChainedCommandHandler(SymbolResultTree symbols, ICommandHandler? commandHandler)
+        {
+            _symbols = symbols;
+            _commandHandler = commandHandler;
+        }
+
+        public int Invoke(InvocationContext context)
+        {
+            // We want to build a stack of (action, next) pairs. But we are not using any collection or LINQ.
+            // Each handler is closure (a lambda with state), where state is the "next" handler.
+            Action<InvocationContext, ICommandHandler?>? chainedHandler = _commandHandler is not null
+                ? (ctx, next) => _commandHandler.Invoke(ctx)
+                : null;
+            ICommandHandler? chainedHandlerArgument = null;
+
+            foreach (var pair in _symbols)
+            {
+                if (pair.Key is Directive directive && directive.HasHandler)
+                {
+                    var syncHandler = directive.SyncHandler 
+                        ?? throw new NotSupportedException($"Directive {directive.Name} does not provide a synchronous handler.");
+
+                    if (chainedHandler is not null)
+                    {
+                        // capture the state in explicit way, to hint the compiler that the current state needs to be used
+                        var capturedHandler = chainedHandler;
+                        var capturedArgument = chainedHandlerArgument;
+
+                        chainedHandlerArgument = new AnonymousCommandHandler(ctx => capturedHandler.Invoke(ctx, capturedArgument));
+                    }
+                    chainedHandler = syncHandler;
+                }
+            }
+
+            chainedHandler!.Invoke(context, chainedHandlerArgument);
+
+            return context.ExitCode;
+        }
+
+        public async Task<int> InvokeAsync(InvocationContext context, CancellationToken cancellationToken = default)
+        {
+            Func<InvocationContext, ICommandHandler?, CancellationToken, Task>? chainedHandler = _commandHandler is not null
+                ? (ctx, next, ct) => _commandHandler.InvokeAsync(ctx, ct)
+                : null;
+            ICommandHandler? chainedHandlerArgument = null;
+
+            foreach (var pair in _symbols)
+            {
+                if (pair.Key is Directive directive && directive.HasHandler)
+                {
+                    var asyncHandler = directive.AsyncHandler
+                        ?? throw new NotSupportedException($"Directive {directive.Name} does not provide an asynchronous handler.");
+
+                    if (chainedHandler is not null)
+                    {
+                        var capturedHandler = chainedHandler;
+                        var capturedArgument = chainedHandlerArgument;
+
+                        chainedHandlerArgument = new AnonymousCommandHandler((ctx, ct) => capturedHandler.Invoke(ctx, capturedArgument, ct));
+                    }
+                    chainedHandler = asyncHandler;
+                }
+            }
+
+            await chainedHandler!.Invoke(context, chainedHandlerArgument, cancellationToken);
+
+            return context.ExitCode;
+        }
+    }
+}

src/System.CommandLine/ParseDirective.cs

@@ -1,6 +1,7 @@
 using System.CommandLine.Invocation;
 using System.CommandLine.IO;
 using System.CommandLine.Parsing;
+using System.Threading.Tasks;
 
 namespace System.CommandLine
 {
@@ -12,17 +13,32 @@ public sealed class ParseDirective : Directive
         /// <param name="errorExitCode">If the parse result contains errors, this exit code will be used when the process exits.</param>
         public ParseDirective(int errorExitCode = 1) : base("parse")
         {
-            SetSynchronousHandler(SyncHandler);
             ErrorExitCode = errorExitCode;
+
+            SetSynchronousHandler(PrintDiagramAndQuit);
+            SetAsynchronousHandler((context, next, cancellationToken) =>
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return Task.FromCanceled(cancellationToken);
+                }
+
+                PrintDiagramAndQuit(context, null);
+
+                return Task.CompletedTask;
+            });
         }
 
         internal int ErrorExitCode { get; }
 
-        private void SyncHandler(InvocationContext context)
+        private void PrintDiagramAndQuit(InvocationContext context, ICommandHandler? next)
         {
             var parseResult = context.ParseResult;
             context.Console.Out.WriteLine(parseResult.Diagram());
             context.ExitCode = parseResult.Errors.Count == 0 ? 0 : ErrorExitCode;
+
+            // parse directive has a precedence over --help and --version and any command
+            // we don't invoke next here.
         }
     }
 }