使用装饰器实现在临时目录中的IO操作

使用python串流程的时候经常会做的操作是：构建所用软件的输入文件-调用软件-读取运行结果-删除产生的文件。如果需要并行的话，可能还需要每个任务单独建一个临时目录，并且报错的时候还希望保留临时目录以方便debug。为了方便，可以使用装饰器自动完成建目录和删除等操作。

def io_in_tempdir(dir='./tmp'):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            temp_dir = tempfile.mkdtemp(dir=dir)
            logger = logging.getLogger("IO")
            logger.debug(f"Tempdir: {temp_dir}")
            try:
                result = func(*args, temp_dir=temp_dir, **kwargs)
            except Exception:
                raise
            else:
                if logging.getLogger().getEffectiveLevel() >= logging.INFO:
                    shutil.rmtree(temp_dir)
                pass
            return result

        return wrapper

    return decorator

以上代码实现了调用时在给定的目录中建立临时目录，返回时删除临时目录，在日志等级低于INFO以及报错的时候会保留临时目录。

在实际串流程的函数中使用subprocess.run获取返回的code，如果不是0就raise error，该错误就会被装饰器捕捉，从而保留临时目录以方便debug。

例如：

@io_in_tempdir()
def prepare_ld(sumstats, ldref, outprefix, temp_dir=None):
    sumstats['SNPID'].to_csv(f'{temp_dir}/snps.txt', index=False, header=False)
    cmd = [
        plink,
        '--bfile',
        ldref,
        '--extract',
        f'{temp_dir}/snps.txt',
        "--r2",
        "square",
        "spaces",
        "--threads",
        "1",
        "--out",
        outprefix,
    ]
    res = run(cmd, stdout=PIPE, stderr=PIPE, universal_newlines=True)
    if res.returncode != 0:
        raise RuntimeError(f"plink failed: {res.stderr}")