doczyai-pipelines/scripts/merge_files.py

#!/usr/bin/env python3
"""
Script to merge multiple text files into a single file.

Usage: 
  python merge_files.py <output_file> <input_file1> <input_file2> ...
  python merge_files.py <output_file> --dir <input_directory> [--pattern <file_pattern>]

This script concatenates the contents of the input files into the output file,
separating each file's content with a newline.
For CSV files, it properly merges by keeping only the header from the first file.
"""

import sys
import os
import glob
import csv

def merge_files(file_list, output_file):
    """
    Merge a list of files into a single output file.

    Args:
        file_list (list): List of file paths to merge.
        output_file (str): Path to the output file.
    """
    if not file_list:
        return

    # Check if files are CSV based on extension
    is_csv = any(fname.lower().endswith('.csv') for fname in file_list)

    if is_csv:
        merge_csv_files(file_list, output_file)
    else:
        merge_text_files(file_list, output_file)

def merge_text_files(file_list, output_file):
    """
    Merge text files by concatenating their contents.
    """
    with open(output_file, 'w', encoding='utf-8') as outfile:
        for fname in file_list:
            if not os.path.isfile(fname):
                print(f"Warning: {fname} is not a file or does not exist. Skipping.")
                continue
            try:
                with open(fname, 'r', encoding='utf-8') as infile:
                    content = infile.read()
                    outfile.write(content)
                    outfile.write('\n')  # Add a newline separator between files
            except Exception as e:
                print(f"Error reading {fname}: {e}")

def merge_csv_files(file_list, output_file):
    """
    Merge CSV files by keeping header from first file and appending data from others.
    """
    first_file = True

    with open(output_file, 'w', newline='', encoding='utf-8') as outfile:
        writer = None

        for fname in file_list:
            if not os.path.isfile(fname):
                print(f"Warning: {fname} is not a file or does not exist. Skipping.")
                continue

            try:
                with open(fname, 'r', encoding='utf-8') as infile:
                    reader = csv.reader(infile)

                    for row_num, row in enumerate(reader):
                        if first_file or row_num > 0:  # Skip header for subsequent files
                            if writer is None:
                                writer = csv.writer(outfile)
                            writer.writerow(row)

                    first_file = False

            except Exception as e:
                print(f"Error reading {fname}: {e}")

def get_files_from_dir(directory, pattern='*'):
    """
    Get all files from a directory matching a pattern.

    Args:
        directory (str): Directory path.
        pattern (str): Glob pattern for files.

    Returns:
        list: List of file paths.
    """
    if not os.path.isdir(directory):
        print(f"Error: {directory} is not a directory.")
        return []
    return glob.glob(os.path.join(directory, pattern))

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python merge_files.py <output_file> <input_file1> <input_file2> ...")
        print("   or: python merge_files.py <output_file> --dir <input_directory> [--pattern <file_pattern>]")
        sys.exit(1)

    output_file = sys.argv[1]

    if sys.argv[2] == '--dir':
        if len(sys.argv) < 4:
            print("Usage: python merge_files.py <output_file> --dir <input_directory> [--pattern <file_pattern>]")
            sys.exit(1)
        directory = sys.argv[3]
        pattern = sys.argv[5] if len(sys.argv) > 5 and sys.argv[4] == '--pattern' else '*'
        input_files = get_files_from_dir(directory, pattern)
    else:
        input_files = sys.argv[2:]

    if not input_files:
        print("No input files found.")
        sys.exit(1)

    merge_files(input_files, output_file)
    print(f"Merged {len(input_files)} files into {output_file}")