让 Python 脚本支持 bash 补全

Question

Bash 有一项很有用的功能，就是当你在终端输入一个命令后,按两下 <TAB> 就会列出该命令支持的所有可能参数。接下来，我会展示如何为一段 python 脚本也赋予这项能力. 而且，该解决方案要足够的通用. 使得他能自动适应对 python 脚本选项或配置文件所做出的改动。

该脚本接受两种类型的参数，一类是以 -- ​开头的标志，一类是从 chef 脚本中取出来的主机名称。

假设该脚本名为 show.py - 它的作用是显示某主机的信息. 它的使用方法为：

  show.py szymon

其中 szymon 为我的 LG 的名字，是从我的 chef 节点定义文件中抽取出来的。

该脚本还支持许多的其他参数，例如：

  show.py --cpu --memory --format=json

脚本支持的参数有两种格式: 一种是简单的字符串，还有一种是以 -- ​开头的参数，为了让该 python 脚本支持 bash 补全，首先需要先写一个简单的 python 脚本来列出所有的 chef node 名称。

  #!/usr/bin/env python

  from sys import argv
  import os
  import json

  if __name__ == "__main__":
      pattern = ""
      if len(argv) == 2:
          pattern = argv[1]

      chef_dir = os.environ.get('LG_CHEF_DIR', None)
      if not chef_dir:
          exit(0)
      node_dirs = [os.path.join(chef_dir, "nodes"),
                   os.path.join(chef_dir, "dev_nodes")]
      node_names = []

      for nodes_dir in node_dirs:
          for root, dirs, files in os.walk(nodes_dir):
              for f in files:
                  try:
                      with open(os.path.join(root, f), 'r') as nf:
                          data = json.load(nf)
                          node_names.append(data['normal']['liquid_galaxy']['support_name'])
                  except:
                      pass

      for name in node_names:
          print name

然后还需要一种方法来获取 python 脚本支持的所有选项. 我们可以通过下面这行代码实现：

  $LG_CHEF_DIR/repo_scripts/show.py --help | grep '  --' | awk {'print $1'}

最后通过一个简单的 bash 脚本来整合上面两个脚本。

  _lg_ssh()
  {
      local cur prev opts node_names
      COMPREPLY=()
      cur="${COMP_WORDS[COMP_CWORD]}"
      prev="${COMP_WORDS[COMP_CWORD-1]}"
      opts=`$LG_CHEF_DIR/repo_scripts/show.py --help | grep '  --' | awk {'print $1'}`
      node_names=`python $LG_CHEF_DIR/repo_scripts/node_names.py`

      if [[ ${cur} == -* ]] ; then
          COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
          return 0
      fi

      COMPREPLY=( $(compgen -W "${node_names}" -- ${cur}) )
  }

  complete -F _lg_ssh show.py

最后，我们只需要在当前 bash session 中使用 source 命令加载该文件即可. 因此，我把下面这行添加到 ~/.bashrc 中。

  source $LG_CHEF_DIR/repo_scripts/show.bash-completion

至此，只要在终端中按两下 <TAB> 就会列出所有的补全选项了。

  $ show.py <tab><tab>
  Display all 42 possibilities? (y or n)
  ... and here go all 42 node names ...
  </tab></tab>

  $ show.py h<tab><tab>
  ... and here go all node names beginning with 'h' ...
  </tab></tab>

  $ show.py --<tab><tab>
  .. and here go all the options beginning with -- ...
  </tab></tab>

该实现最大的好处在于，任何对 python script 的修改或 chef node 名称的修改，都会自动反映出来。