aws-vaultでexecするprofileとコマンドをzshの補完で選べるようにする

aws-vaultはAWSのprofileを指定して任意のコマンドを実行できます。

$ aws-vault exec my-profile -- aws s3 ls

この profile(my-profile)の部分と 任意のコマンド(aws s3 ls)の部分をTabキーで補完できるようにしました。

zsh

# サブコマンドの補完
$ aws-vault [Tab]
Completing subcommand
add     clear   exec    help    list    login   remove  rotate

# profileの補完
$ aws-vault exec [Tab]
Completing profile
bar-role1     bar-role2     default       foo-admin     foo-readonly  jonsmith

$ aws-vault exec f[Tab]
$ aws-vault exec foo-[Tab]
Completing profile
foo-admin     foo-readonly

# コマンドの補完
$ aws-vault exec foo-readonly [Tab]
$ aws-vault exec foo-readonly -- [Tab]
Completing external command
.keepme                             ifconfig                            pnmtopng                          
2to3                                ifnames                             pnmtopnm                          
2to3-                               ilbmtoppm                           pnmtops                           
2to3-2.7                            img2webp                            pnmtorast                         
2to3-3.9                            imgcmp                              pnmtorle
...

$ aws-vault exec foo-readonly -- aws e[Tab]
ebs                   ecr-public            elastic-inference     elb                   es                                        
ec2                   ecs                   elasticache           elbv2                 events                                    
ec2-instance-connect  efs                   elasticbeanstalk      emr                   evidently                                 
ecr                   eks                   elastictranscoder     emr-containers

ソースコード

https://github.com/naitoNanaco/aws-vault-completion

解説

補完の基本

補完スクリプト名ファイルを用意し、

• 1行目に #compdef コマンド名
• 補完スクリプト名関数
• 最終行にcompdef 補完スクリプト名 コマンド名
  を記述します。

補完スクリプトの中で、
_values関数を使うと、第1引数が補完グループ名、第２引数以降が補完の候補として設定されます。
_arguments関数を使うと、引数の位置に応じたアクションを定義できます。

サブコマンドの補完

_arguments関数のアクションとして->(state-val)と記述することでstate変数に値がセットされます。
aws-vaultの第１引数を処理する際は$stateに"subcommand"をセットし、定義したサブコマンドのリストを補完の候補とします。

_aws-vault

#compdef aws-vault

_aws-vault() {
  # サブコマンドのリスト（補完の候補）
  local -a val
  val=(help add list rotate exec clear remove login)

  _arguments \
    '1: :->subcommand' \    # 1つ目の引数処理時のstateを"subcommand"としてセット

  case "$state" in
    subcommand)
      _values $state $val   # stateが"subcommand"の際はサブコマンドのリストが候補とする
      ;;
  esac
}

compdef _aws-vault aws-vault

profileの補完

第２引数以降はサブコマンドに何を選んだかに応じて以降の補完アクションを分岐させます。
第２引数以降の処理時、_arguments関数で$stateに"args"をセットします。
$stateが"args"だった場合に、$words[1]の値を確認することで第１引数がどれか確認できます。
execだった場合に、改めて_arguments関数を用いて初めの引数処理時は__aws-vault-profile関数が呼ばれるようにします。
__aws-vault-profile関数ではprofileのリストをaws-vault lsの結果から取得し、_value関数によって補完の候補に設定しています。

_aws-vault

_aws-vault() {

  _arguments \
    '1: :->subcommand' \    # 1つ目の引数処理時のstateを"subcommand"としてセット
    '*:: :->args' \         # 以降の引数処理時のstateを"args"としてセット

  case "$state" in
    subcommand)
      ;;
    args)
      # サブコマンドに何を選んだかに応じて以降の補完アクションを分岐（今回はexecのみ定義）
      case $words[1] in
        exec)
          _arguments \
            '1: :__aws-vault-profile' \  # 1つ目の引数処理時に"__aws-vault-profile"関数を実行
            '*:: :__aws-vault-exec-precommand' \
          ;;
      esac
      ;;
  esac
}

__aws-vault-profile() {
  # profileのリストを`aws-vault ls`の結果から取得
  local -a val
  val=(`aws-vault ls | awk 'NR>2 {if ($1 != "-") print $1}'`)
  # 取得したリストを補完の候補とする
  _values 'profile' $val
}

コマンドの補完

profileの選択以降は__aws-vault-exec-precommand関数で処理しています。
1つめの処理の際には区切り文字--を補完するように__aws-vault-delimiter関数を呼び出します。
それ以降はzshに定義された他の補完に全てを委任します。委任には_command関数を使います。
（参考：zshに完全な自動補完を継承させるにはどうすればよいですか? - CODE Q&A）
__aws-vault-exec-command関数内で_command関数を呼び出します。

_aws-vault

_aws-vault() {

  _arguments \
    '1: :->subcommand' \
    '*:: :->args' \

  case "$state" in
    subcommand)
      ;;
    args)
      case $words[1] in
        exec)
          _arguments \
            '1: :__aws-vault-profile' \
            '*:: :__aws-vault-exec-precommand' \
          ;;
      esac
      ;;
  esac
}

__aws-vault-exec-precommand() {
  _arguments \
    '1: :__aws-vault-delimiter' \
    '*:: :__aws-vault-exec-command' \
}

__aws-vault-delimiter() {
  local -a val
  val=(--)
  _values 'EOC' $val
}

__aws-vault-exec-command() {
  _command aws-vault
}

参考

参考にさせていただきました。

• Zsh: 自作コマンドで補完できるようにする
• zsh の補完関数の自作導入編
• zshの補完関数を作ってみた