CSV文件合并¶

详情参见 GitHub
'''
CSV文件合并

作者：HuangFuSL
'''

import argparse
import os
import itertools
import csv
from typing import Iterator, List, Tuple, TypeVar

T = TypeVar('T')

def extract_head(s: Iterator[T]) -> Tuple[T, List[T]]:
    return next(s), list(s)

def load_csv(file: str):
    with open(file, 'r', encoding='utf-8', newline='') as f:
        dialect = csv.Sniffer().sniff(f.read(), delimiters=',')
        f.seek(0)
        return extract_head(csv.reader(f, dialect=dialect))


def merge_csv(dest: str, *src: str, encoding: str):
    dest_head = tuple()
    wf = open(dest, 'w', encoding=encoding, newline='')
    writer = csv.writer(wf)
    for file in src:
        head, content = load_csv(file)
        if not dest_head:
            dest_head = tuple(head)
            writer.writerow(dest_head)

        writer.writerows(content)
    wf.close()

def search(path: str, recursive: bool):
    if not recursive:
        yield from [
            os.path.join(path, f)
            for f in os.listdir(path)
                if f.endswith('.csv')
        ]
    for root, _, files in os.walk(path):
        yield from [
            os.path.join(path, root, f)
            for f in files
                if f.endswith('.csv')
        ]

parser = argparse.ArgumentParser(description='Merge CSV Sheets')
parser.add_argument(
    '-o', '--output', nargs=1, default='output.csv',
    help='Path to output file'
    )
parser.add_argument(
    '-R', '--recursive', action='store_true',
    help='Search files recursively when directory is specified'
)
parser.add_argument(
    '-e', '--encoding', default='utf-8',
    help='Output encoding'
)
parser.add_argument('files', nargs='*', help='Files to merge')

if __name__ == '__main__':
    ns = parser.parse_args()
    files = [_ for _ in ns.files if os.path.isfile(_)]
    dirs = [_ for _ in ns.files if os.path.isdir(_)]
    targets = itertools.chain(files, *map(lambda d: search(d, ns.recursive), dirs))
    merge_csv(ns.output, *targets, encoding=ns.encoding)
CSV文件合并¶

评论
