add WithContext func to gopsutil

contextnize.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import

# python contextanize.py *.go
# after that, goimports -w is required to add import "context"
#
# ex
# +       return TimesWithContext(context.Background())
# +}
# +
# +func TimesWithContext(ctx context.Context, percpu bool) ([]TimesStat, error) {


import sys
import re

funcre = re.compile("func (.+?)\((.*){")
procre = re.compile("func \(p \*Process\) (.+?)\((.*?){")

def print_context_process(line: str) -> list:
    buf = [line]

    m = re.match(funcre, line.replace("(p *Process) ", ""))
    if m is None:
        return buf

    fn = m.group(1)

    args = m.groups(2)[1].split(")")[0]  # get argments
    args = [arg.strip().split(" ")[0] for arg in args.split(",")]
    if args[0] == "":
        buf.append(f"\treturn p.{fn}WithContext(context.Background())\n")
    else:
        args = ", ".join(args)
        buf.append(f"\treturn p.{fn}WithContext(context.Background(), {args})\n")
    buf.append("}\n")
    buf.append("\n")
    buf.append(f"func (p *Process) {fn}WithContext(ctx context.Context, {m.group(2)}{{\n")

    return buf

def print_context(line: str) -> list:
    if "*Process" in line:
        return print_context_process(line)

    buf = [line]

    m = re.match(funcre, line)
    if m is None:
        return buf

    fn = m.group(1)

    args = m.groups(2)[1].split(")")[0]  # get argments
    args = [arg.strip().split(" ")[0] for arg in args.split(",")]
    args = ", ".join(args)

    if "..." in m.groups(2)[1].split(")")[0]:
        buf.append(f"\treturn {fn}WithContext(context.Background(), {args}...)\n")
    else:
        buf.append(f"\treturn {fn}WithContext(context.Background(), {args})\n")
    buf.append("}\n")
    buf.append("\n")
    buf.append(f"func {fn}WithContext(ctx context.Context, {m.group(2)}{{\n")

    return buf

def add(filename: str):
    buf = []

    blacklist = [
        "testing",
        "String()",
        "NewProcess",
        "Processes",
    ]
    def check(line):
        for b in blacklist:
            if b in line:
                return True

    with open(filename) as f:
        for line in f:
            if not line.startswith("func "):
                buf.append(line)
                continue
            if check(line):
                buf.append(line)
                continue
            if line[5] == "(" and "*Process" not in line:
                buf.append(line)
                continue
            if line[5].islower():
                buf.append(line)
                continue
            buf += print_context(line)

    with open(filename, "w") as f:
        for b in buf:
            f.write(b)


for name in sys.argv[1:]:
    if not name.endswith(".go"):
        continue
    if name.endswith("_test.go"):
        continue
    add(name)