54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
import csv
|
|
import os
|
|
from django.core.management.base import BaseCommand
|
|
|
|
class Command(BaseCommand):
|
|
help = 'Flags staff as heads based on whether they appear in the Manager column'
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument('input_csv', type=str, help='Path to original CSV')
|
|
parser.add_argument('output_csv', type=str, help='Path to save the new CSV')
|
|
|
|
def handle(self, *args, **options):
|
|
input_path = options['input_csv']
|
|
output_path = options['output_csv']
|
|
|
|
if not os.path.exists(input_path):
|
|
self.stdout.write(self.style.ERROR(f"File {input_path} not found."))
|
|
return
|
|
|
|
rows = []
|
|
manager_ids_set = set()
|
|
|
|
# Pass 1: Collect all IDs from the Manager column
|
|
with open(input_path, 'r', encoding='utf-8') as f:
|
|
reader = csv.DictReader(f)
|
|
fieldnames = reader.fieldnames
|
|
|
|
for row in reader:
|
|
rows.append(row)
|
|
manager_val = row.get('Manager', '')
|
|
|
|
if manager_val and '-' in manager_val:
|
|
# Split by '-' and take the first part (the ID)
|
|
m_id = manager_val.split('-')[0].strip()
|
|
if m_id:
|
|
manager_ids_set.add(m_id)
|
|
|
|
# Pass 2: Flag rows as 'is_head' if their Staff ID is in the set
|
|
for row in rows:
|
|
staff_id = row.get('Staff ID', '').strip()
|
|
if staff_id in manager_ids_set:
|
|
row['is_head'] = 1
|
|
else:
|
|
row['is_head'] = 0
|
|
|
|
# Pass 3: Write the new CSV
|
|
new_fieldnames = fieldnames + ['is_head']
|
|
with open(output_path, 'w', encoding='utf-8', newline='') as f:
|
|
writer = csv.DictWriter(f, fieldnames=new_fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
self.stdout.write(self.style.SUCCESS(f"Processed {len(rows)} rows. Found {len(manager_ids_set)} unique heads."))
|
|
self.stdout.write(self.style.SUCCESS(f"Saved to: {output_path}")) |