Prints PIDs for given program basename (or wildcard if not found)

pids.py

#!/usr/bin/env python3

from collections import defaultdict
import os
import sys

def read_cmdline(file):
    try:
        with open(file, "rb") as f:
            cmd = str(f.read(), encoding="utf-8")
            if cmd is not None:
                return cmd.replace("\x00", " ")
    except:
        pass

def getprogs():
    # Index on the basename of cmdline
    basenames = defaultdict(lambda: [])
    # Index on the full command line
    fullnames = defaultdict(lambda: [])

    for pid in os.listdir("/proc/"):
        if pid.isdigit():
            pid = int(pid)
            cmd = read_cmdline("/proc/%d/cmdline" % pid)
            if cmd is not None:
                basenames[os.path.basename(cmd.split(" ")[0])].append(pid)
                fullnames[cmd].append(pid)
    return basenames, fullnames

if __name__ == "__main__":
    if ("-h" in sys.argv) or ("--help" in sys.argv):
        print("Usage: pid name [ name]")
        print("Returns pids of program basenames.")
        print("Example: pid vim # prints pids for all vim procs")
        sys.exit(0)

    self = os.getpid()

    basenames, fullnames = getprogs()
    for prog in sys.argv[1:]:
        if prog in basenames:
            for pid in basenames[prog]:
                if pid != self:
                    print(pid)
        else:
            for cmd, pids in fullnames.items():
                if prog in cmd:
                    for pid in pids:
                        if pid != self:
                            print(pid)