102 lines
3.3 KiB
Python
Executable file
102 lines
3.3 KiB
Python
Executable file
#!/usr/bin/env python
|
|
import csv
|
|
import sys
|
|
import os
|
|
from datetime import datetime
|
|
|
|
# Get the project root directory
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.dirname(script_dir)
|
|
data_dir = os.path.join(project_root, 'data')
|
|
|
|
def get_output_filename(input_filename, chunk_number):
|
|
base, ext = os.path.splitext(os.path.basename(input_filename))
|
|
return os.path.join(data_dir, f"{base}-chunk-{chunk_number}{ext}")
|
|
|
|
def get_history_filename(input_filename):
|
|
base, ext = os.path.splitext(os.path.basename(input_filename))
|
|
return os.path.join(data_dir, f"{base}-history{ext}")
|
|
|
|
def read_csv(filename):
|
|
with open(filename, 'r', newline='') as csvfile:
|
|
reader = csv.reader(csvfile)
|
|
return list(reader)
|
|
|
|
def write_csv(filename, data):
|
|
with open(filename, 'w', newline='') as csvfile:
|
|
writer = csv.writer(csvfile)
|
|
writer.writerows(data)
|
|
|
|
def update_history(history_filename, new_data, chunk_number):
|
|
if os.path.exists(history_filename):
|
|
history_data = read_csv(history_filename)
|
|
headers = history_data[0]
|
|
if 'extraction' not in headers:
|
|
headers.append('extraction')
|
|
for row in history_data[1:]:
|
|
row.append('chunk-1')
|
|
else:
|
|
headers = new_data[0] + ['extraction']
|
|
history_data = [headers]
|
|
|
|
for row in new_data[1:]:
|
|
history_data.append(row + [f'chunk-{chunk_number}'])
|
|
|
|
write_csv(history_filename, history_data)
|
|
|
|
def get_extracted_rows(history_filename):
|
|
if not os.path.exists(history_filename):
|
|
return set()
|
|
|
|
history_data = read_csv(history_filename)
|
|
extracted_rows = set()
|
|
for row in history_data[1:]: # Skip header
|
|
extracted_rows.add(tuple(row[:-1])) # Exclude the 'extraction' column
|
|
return extracted_rows
|
|
|
|
def extract_targets(input_filename, num_rows):
|
|
full_input_path = os.path.join(data_dir, input_filename)
|
|
history_filename = get_history_filename(input_filename)
|
|
extracted_rows = get_extracted_rows(history_filename)
|
|
|
|
input_data = read_csv(full_input_path)
|
|
headers = input_data[0]
|
|
|
|
new_chunk = [headers]
|
|
chunk_number = len(set(row[-1] for row in read_csv(history_filename)[1:])) + 1 if os.path.exists(history_filename) else 1
|
|
|
|
rows_added = 0
|
|
for row in input_data[1:]:
|
|
if tuple(row) not in extracted_rows and rows_added < num_rows:
|
|
new_chunk.append(row)
|
|
rows_added += 1
|
|
|
|
if rows_added == num_rows:
|
|
break
|
|
|
|
if rows_added == 0:
|
|
print("No new rows to extract.")
|
|
return
|
|
|
|
output_filename = get_output_filename(input_filename, chunk_number)
|
|
write_csv(output_filename, new_chunk)
|
|
update_history(history_filename, new_chunk, chunk_number)
|
|
|
|
print(f"Extracted {rows_added} rows to {output_filename}")
|
|
print(f"Updated history file: {history_filename}")
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 3:
|
|
print("Usage: python extract-targets.py <input_filename> <num_rows>")
|
|
sys.exit(1)
|
|
|
|
input_filename = sys.argv[1]
|
|
num_rows = int(sys.argv[2])
|
|
|
|
full_input_path = os.path.join(data_dir, input_filename)
|
|
if not os.path.exists(full_input_path):
|
|
print(f"Error: Input file '{full_input_path}' not found.")
|
|
sys.exit(1)
|
|
|
|
extract_targets(input_filename, num_rows)
|