domain-scraping-utils/bin/split-inputs.py
2024-07-18 15:52:01 -04:00

63 lines
1.9 KiB
Python
Executable file

#!/usr/bin/env python
import os
import csv
import math
import sys
# Increase CSV field size limit
csv.field_size_limit(sys.maxsize)
# Determine the project root dynamically
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
INPUT_FILE = os.path.join(PROJECT_ROOT, "input-domains.csv")
OUTPUT_FILES = [os.path.join(PROJECT_ROOT, f"input-domains-{i}.csv") for i in range(1, 4)]
def split_targets():
print(f"Reading input file: {INPUT_FILE}")
if not os.path.exists(INPUT_FILE):
print(f"Error: Input file {INPUT_FILE} does not exist.")
sys.exit(1)
# Read all rows from the input file
try:
with open(INPUT_FILE, 'r', encoding='utf-8', errors='replace') as f:
reader = csv.reader(f)
headers = next(reader) # Read the header row
rows = list(reader)
except Exception as e:
print(f"Error reading input file: {e}")
sys.exit(1)
total_rows = len(rows)
print(f"Total rows: {total_rows}")
if total_rows == 0:
print("No rows found. Exiting.")
sys.exit(0)
# Calculate split sizes
base_size = math.ceil(total_rows / 3)
split_sizes = [base_size] * 2 + [total_rows - base_size * 2]
# Write split files
start_index = 0
for i, output_file in enumerate(OUTPUT_FILES):
end_index = start_index + split_sizes[i]
try:
with open(output_file, 'w', newline='', encoding='utf-8') as f:
writer = csv.writer(f)
writer.writerow(headers)
writer.writerows(rows[start_index:end_index])
print(f"Written {split_sizes[i]} rows to {output_file}")
except Exception as e:
print(f"Error writing to {output_file}: {e}")
start_index = end_index
def main():
print("Starting split-targets.py")
split_targets()
print("split-targets.py completed")
if __name__ == "__main__":
main()