csv_headers.py

# -*- coding: utf-8 -*-
import argparse
import gzip
import logging
import logging.config
import os
import sys
import zipfile

logger = logging.getLogger(__name__)


def open_report(cur_report):
    '''
        Open report and output all the section headers.
    '''

    if not os.path.isfile(cur_report):
        logger.critical("File not found: {}".format(cur_report))
        sys.exit(1)

    _, compression_type = cur_report.rsplit(".", 1)
    report_path, base_report = cur_report.rsplit(r'/', 1)
    source_csv = None

    if compression_type.lower() == 'zip':
        z = zipfile.ZipFile(cur_report)
        internal_name = z.namelist()[0]

        if internal_name != base_report.rsplit(".", 1)[0]:
            logger.info("Warning unexpected name: {}". format(internal_name))

        source_csv = z.open(internal_name)
    elif compression_type.lower() == 'gz':
        source_csv = gzip.open(cur_report, 'rb')
    elif compression_type.lower() == 'csv':
        source_csv = open(cur_report, 'rb')
    else:
        logger.critical("Unsupported extension: {}".format(cur_report))
        sys.exit(2)

    def clean_header(row):
        return row.replace(' ', '_').replace('-', '_').lower()

    header = True
    first = True
    for line in source_csv:
        if first is True and len(line.split(',')) > 1:
            first = False
            logger.info(clean_header(line))
            logger.debug("Single Section Detected")
            break

        if header is True:
            header, first = False, False
            logger.info(clean_header(line))

        # Empty Section
        if len(line) <= 1:
            header = True

    source_csv.close()

if __name__ == "__main__":
    '''
        Print out all the section headers from a YouTube report.
    '''
    log_std = logging.StreamHandler(sys.stdout)
    log_std.setFormatter(logging.Formatter('%(message)s'))
    log_std.setLevel(logging.DEBUG)
    logger.addHandler(log_std)
    logger.setLevel(logging.INFO)

    parser = argparse.ArgumentParser(description='Get Headers from csv files.')
    parser.add_argument(dest='report_files', metavar='file', nargs='+',
                        help='Youtube Report for header extraction.')

    args = parser.parse_args()

    for report_file in args.report_files:
        logger.info("Headers for {}".format(report_file))
        open_report(report_file)