From 80f2db8c15dbac0662ce2b74edf4ec6813f04d47 Mon Sep 17 00:00:00 2001
From: Konrad Baumgart <konrad.baumgart@erli.pl>
Date: Fri, 10 Dec 2021 08:36:35 +0100
Subject: [PATCH] completion for positional parameters with options

---
 lib/completion.ts   |  23 +++++++++
 test/completion.cjs | 123 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 144 insertions(+), 2 deletions(-)

diff --git a/lib/completion.ts b/lib/completion.ts
index 60cc486e0..5a591d273 100644
--- a/lib/completion.ts
+++ b/lib/completion.ts
@@ -32,6 +32,7 @@ export class Completion implements CompletionInstance {
 
   private aliases: DetailedArguments['aliases'] | null = null;
   private customCompletionFunction: CompletionFunction | null = null;
+  private indexAfterLastReset = 0;
   private readonly zshShell: boolean;
 
   constructor(
@@ -57,6 +58,7 @@ export class Completion implements CompletionInstance {
       if (handlers[args[i]] && handlers[args[i]].builder) {
         const builder = handlers[args[i]].builder;
         if (isCommandBuilderCallback(builder)) {
+          this.indexAfterLastReset = i + 1;
           const y = this.yargs.getInternalMethods().reset();
           builder(y, true);
           return y.argv;
@@ -146,6 +148,27 @@ export class Completion implements CompletionInstance {
         completions.push(...choices);
       }
     }
+
+    const positionalKeys =
+      this.yargs.getGroups()[this.usage.getPositionalGroupName()] || [];
+    const offset = Math.max(
+      this.indexAfterLastReset,
+      this.yargs.getInternalMethods().getContext().commands.length +
+        /* name of the script is first param */ 1
+    );
+    const positionalValues = argv._.slice(offset);
+
+    const positionalKey = positionalKeys[positionalValues.length - 1];
+
+    if (!positionalKey) {
+      return;
+    }
+    const choices = this.yargs.getOptions().choices[positionalKey] || [];
+    for (const choice of choices) {
+      if (choice.startsWith(current)) {
+        completions.push(choice);
+      }
+    }
   }
 
   private getPreviousArgChoices(args: string[]): string[] | void {
diff --git a/test/completion.cjs b/test/completion.cjs
index d4b934178..323b7aba3 100644
--- a/test/completion.cjs
+++ b/test/completion.cjs
@@ -17,8 +17,8 @@ describe('Completion', () => {
 
   describe('default completion behavior', () => {
     const firstArgumentOptions = [
-      ['--get-yargs-completions'], // proper args according to docs
-      ['./completion', '--get-yargs-completions'], // yargs is called like this in tests, don't know why
+      ['--get-yargs-completions'], // proper args after hideBin is used
+      ['./completion', '--get-yargs-completions'], // yargs is called like this in tests a lot
     ];
     for (const firstArguments of firstArgumentOptions) {
       describe(`calling yargs(${firstArguments.join(', ')}, …)`, () => {
@@ -333,6 +333,125 @@ describe('Completion', () => {
           r.logs.should.not.include('banana');
           r.logs.should.not.include('pear');
         });
+
+        it('completes choices for first positional', () => {
+          process.env.SHELL = '/bin/bash';
+          const r = checkUsage(
+            () =>
+              yargs([...firstArguments, './completion', 'cmd', ''])
+                .help(false)
+                .version(false)
+                .command('cmd [fruit] [fruit2]', 'command', subYargs => {
+                  subYargs
+                    .positional('fruit', {choices: ['apple', 'banana', 'pear']})
+                    .positional('fruit2', {
+                      choices: ['apple2', 'banana2', 'pear2'],
+                    })
+                    .options({amount: {describe: 'amount', type: 'number'}});
+                }).argv
+          );
+
+          r.logs.should.have.length(4);
+          r.logs.should.include('apple');
+          r.logs.should.include('banana');
+          r.logs.should.include('pear');
+          r.logs.should.include('--amount');
+        });
+
+        it('completes choices for positional with prefix', () => {
+          process.env.SHELL = '/bin/bash';
+          const r = checkUsage(
+            () =>
+              yargs([...firstArguments, './completion', 'cmd', 'a'])
+                .help(false)
+                .version(false)
+                .command('cmd [fruit] [fruit2]', 'command', subYargs => {
+                  subYargs
+                    .positional('fruit', {
+                      choices: ['apple1', 'banana1', 'pear1'],
+                    })
+                    .positional('fruit2', {
+                      choices: ['apple2', 'banana2', 'pear2'],
+                    })
+                    .options({amount: {describe: 'amount', type: 'number'}});
+                }).argv
+          );
+
+          r.logs.should.have.length(1);
+          r.logs.should.include('apple1');
+        });
+
+        it('completes choices for second positional after option', () => {
+          process.env.SHELL = '/bin/bash';
+          const r = checkUsage(
+            () =>
+              yargs([
+                ...firstArguments,
+                './completion',
+                'cmd',
+                'apple',
+                '--amount',
+                '1',
+                '',
+              ])
+                .help(false)
+                .version(false)
+                .command('cmd [fruit] [fruit2]', 'command', subYargs => {
+                  subYargs
+                    .positional('fruit', {
+                      choices: ['apple1', 'banana1', 'pear1'],
+                    })
+                    .positional('fruit2', {
+                      choices: ['apple2', 'banana2', 'pear2'],
+                    })
+                    .options({amount: {describe: 'amount', type: 'number'}});
+                }).argv
+          );
+
+          r.logs.should.have.length(3);
+          r.logs.should.include('apple2');
+          r.logs.should.include('banana2');
+          r.logs.should.include('pear2');
+        });
+
+        it('completes choices for nested command', () => {
+          process.env.SHELL = '/bin/bash';
+          const r = checkUsage(
+            () =>
+              yargs([...firstArguments, './completion', 'wrapper', 'cmd', ''])
+                .help(false)
+                .version(false)
+                .command({
+                  command: 'wrapper',
+                  builder: subYargs =>
+                    subYargs.command('cmd [fruit]', 'command', subYargs2 => {
+                      subYargs2.positional('fruit', {choices: ['apple']});
+                    }),
+                }).argv
+          );
+
+          r.logs.should.have.length(1);
+          r.logs.should.include('apple');
+        });
+
+        it('works if positional has no choices', () => {
+          process.env.SHELL = '/bin/bash';
+          const r = checkUsage(
+            () =>
+              yargs([...firstArguments, './completion', 'wrapper', 'cmd', 'a'])
+                .help(false)
+                .version(false)
+                .command({
+                  command: 'wrapper',
+                  builder: subYargs =>
+                    subYargs.command('cmd [fruit]', 'command', subYargs2 => {
+                      subYargs2.positional('fruit', {});
+                    }),
+                }).argv
+          );
+
+          r.logs.should.have.length(0);
+        });
       });
     }
   });