Source code for unravel.tabular.columns

#!/usr/bin/env python3

"""
Use ``tabular_columns`` (``cols``) from UNRAVEL to load a CSV or xlsx file and print the column(s).

Usage:
------
    tabular_columns --input path/input.csv [-one-per-line] [--rows N] [-d delimiter] [-v]
"""

import pandas as pd
from rich import print
from rich.traceback import install

from unravel.core.help_formatter import RichArgumentParser, SuppressMetavar, SM
from unravel.core.config import Configuration 
from unravel.core.utils import log_command, verbose_start_msg, verbose_end_msg




def parse_args():
    parser = RichArgumentParser(formatter_class=SuppressMetavar, add_help=False, docstring=__doc__)

    reqs = parser.add_argument_group('Required arguments')
    reqs.add_argument('-i', '--input', required=True, help='Path to the input CSV or xlsx file.', action=SM)

    opts = parser.add_argument_group('Optional arguments')
    opts.add_argument('-1', '--one-per-line', help='Print each column name on a separate line.', action='store_true', default=False)
    opts.add_argument('-r', '--rows', help='Print the first N rows. Default: 0 (only print column names)', type=int, default=0, action=SM)
    opts.add_argument('-d', '--delimiter', help='Delimiter used in the CSV file. Default: ", "', default=', ')

    general = parser.add_argument_group('General arguments')
    general.add_argument('-v', '--verbose', help='Increase verbosity. Default: False', action='store_true', default=False)

    return parser.parse_args()





@log_command
def main():
    install()
    args = parse_args()
    Configuration.verbose = args.verbose
    verbose_start_msg()

    # Load the CSV header
    if str(args.input).endswith('.csv'):
        cols = pd.read_csv(args.input, nrows=0).columns
    elif str(args.input).endswith('.xlsx'):
        cols = pd.read_excel(args.input, nrows=0).columns

    # Print column names
    if args.one_per_line:
        for col in cols:
            if args.rows > 0:
                print(f'[green]{col}[/green][default]{args.delimiter}[/default]{args.delimiter.join(map(str, pd.read_csv(args.input, usecols=[col], nrows=args.rows).values.flatten()))}')
            else:
                print(f'[default]{col}')
    else:
        if args.rows > 0:
            header = f'[default]{args.delimiter}[/default]'.join(
                f'[green]{col}[/green]' for col in cols
            )
            values = args.delimiter.join(
                map(str, pd.read_csv(args.input, usecols=cols, nrows=args.rows).values.flatten())
            )
            print(f'{header}[default]{args.delimiter}[/default]{values}')        
        else:
            print("Columns in the CSV file:")
            print(f"[default]{args.delimiter}".join(cols))

    verbose_end_msg()



if __name__ == '__main__':
    main()