コマンドラインオプション付きの which コマンドを実装してみた

python

というツイートを見かけました。

以前、勉強がてらに

を作ったことがありましたが、もうあれから3年も経つんだなとコードを見返しながら感慨に浸っていました。

そんなとき、ふと which コマンドを実行したところ、コマンドラインオプションがあることに気付きました！

$ which
usage: which [-as] program ...

$ man which
WHICH(1)                  BSD General Commands Manual                 WHICH(1)

NAME
     which -- locate a program file in the user's path

SYNOPSIS
     which [-as] program ...

DESCRIPTION
     The which utility takes a list of command names and searches the path for each 
     executable file that would be run had these commands actually been invoked.

     The following options are available:

     -a      List all instances of executables found (instead of just the first one of each).

     -s      No output, just return 0 if any of the executables are found, or 1 if none are found.

     Some shells may provide a builtin which command which is similar or identical
     to this utility.  Consult the builtin(1) manual page.

これは、、、

。。。

過去の私の実装に足りないものを発見してしまいました、もう再実装するしかない！ (> <)

ということで、コマンドラインオプションを追加するだけなのですが、結構難しくて悩んでしまいました。

# -*- coding: utf-8 -*-
import glob
import os
import sys
import argparse
from itertools import chain
from os.path import join as pathjoin
from operator import itemgetter

def search(cmd, paths, is_all=False):
    for path in paths:
        for match in glob.glob(pathjoin(path, cmd)):
            if os.access(match, os.X_OK):
                yield match
                if not is_all:
                    raise StopIteration

def parse_argument(args=None):
    parser = argparse.ArgumentParser()
    parser.set_defaults(is_all=False, is_silent=False, commands=[])
    parser.add_argument("-a", dest="is_all", action="store_true",
        help="List all instances of executables found "
             "(instead of just the first one of each).")
    parser.add_argument("-s", dest="is_silent", action="store_true",
        help="No output, just return 0 if any of the executables are found, "
             "or 1 if none are found.")
    parser.add_argument("commands", nargs="*")
    args = parser.parse_args(args or sys.argv[1:])
    return args

def main(cmd_args=None):
    args = parse_argument(cmd_args)
    env_paths = os.environ['PATH'].split(':')
    result = []
    for cmd in args.commands:
        founds = list(search(cmd, env_paths, args.is_all))
        result.append((0, founds) if founds else (1, [cmd]))

    status_code = max(map(itemgetter(0), result))
    if not args.is_silent:
        cmd_paths = [paths for ret_val, paths in result if ret_val == 0]
        for cmd_path in chain.from_iterable(cmd_paths):
            print cmd_path
    return status_code

if __name__ == '__main__':
    sys.exit(main())

実行結果はこんな感じです。

$ python which.py -a ls vi unknown
/bin/ls
/opt/local/bin/vi
/usr/bin/vi

$ python which.py -s ls vi unknown; echo $?
1

ついでに本当に which コマンドと動作が一致しているか、簡単なテストを書いてみました。

# -*- coding: utf-8 -*-
import sys
from subprocess import Popen, PIPE

import pytest
import which

FILESYSTEM_ENCODING = sys.getfilesystemencoding()

def get_system_which_result(args):
    cmds = ["which"]
    cmds.extend(args)
    p = Popen(" ".join(cmds), stdout=PIPE, stderr=PIPE, shell=True)
    out, err = p.communicate()
    return p.returncode, out, err

@pytest.mark.parametrize("args", [
    ["ls"],
    ["cd", "pwd"],
    ["non_existence"],
    ["-a", "vi"],
    ["-a", "ls", "vi"],
    ["-s", "non_existence"],
    ["-a", "-s", "ls", "vi"],
    ["-a", "-s", "ls", "vi", "non_existence"],
])
def test_which_command(args, capsys):
    my_ret = which.main(args)
    my_out, my_err = capsys.readouterr()
    sys_ret, sys_out, sys_err = get_system_which_result(args)
    assert sys_ret == my_ret
    assert sys_out == my_out.encode(FILESYSTEM_ENCODING)
    assert sys_err == my_err.encode(FILESYSTEM_ENCODING)

ソースは以下にあります。

こういったシンプルなツールを作るのはおもしろいですね。