"""Bulk-add dataset collaborators to a Segments.ai dataset from a CSV of usernames and roles. The CSV is expected to have "username" and "role" columns, e.g.: username,role john,reviewer jane,labeler Valid roles: labeler, reviewer, manager, admin. Rows with an empty role default to "labeler". """ import csv import json from segments import SegmentsClient from segments.exceptions import SegmentsError VALID_ROLES = ["labeler", "reviewer", "manager", "admin"] DEFAULT_ROLE = "labeler" def format_error_message(error: SegmentsError) -> str: body, separator, possible_fixes = str(error).partition("\nPossible fixes:") try: parsed_body = json.loads(body.strip()) except json.JSONDecodeError: parsed_body = None if isinstance(parsed_body, dict): body = " ".join(str(item) for value in parsed_body.values() for item in (value if isinstance(value, list) else [value])) return f"{body.strip()}{separator}{possible_fixes}".strip() def read_collaborators(csv_path: str) -> tuple[list[tuple[str, str]], list[str]]: with open(csv_path, newline="") as csv_file: reader = csv.DictReader(csv_file) if reader.fieldnames is None or "username" not in reader.fieldnames: raise ValueError('CSV must have a "username" column') collaborators = [] skipped_rows = [] for row_index, row in enumerate(reader, start=2): username = row["username"].strip() if username == "": skipped_rows.append(f"Row {row_index}: skipped, missing username") continue role = (row.get("role") or "").strip() or DEFAULT_ROLE if role not in VALID_ROLES: raise ValueError(f"Row {row_index}: invalid role '{role}' for username '{username}'. Must be one of {VALID_ROLES}") collaborators.append((username, role)) return collaborators, skipped_rows def add_collaborators(client: SegmentsClient, dataset_identifier: str, collaborators: list[tuple[str, str]]) -> tuple[list[str], bool]: logs = [] failed_usernames = [] usernames_by_error_message: dict[str, list[str]] = {} for username, role in collaborators: try: client.add_dataset_collaborator(dataset_identifier, username, role) logs.append(f"Added {username} to {dataset_identifier} as {role}") except SegmentsError as error: message = format_error_message(error) usernames_by_error_message.setdefault(message, []).append(username) failed_usernames.append(username) for message, usernames in usernames_by_error_message.items(): logs.append(f"Failed to add {', '.join(usernames)}: {message}\n") if failed_usernames: logs.append(f"{len(failed_usernames)} of {len(collaborators)} usernames failed: {', '.join(failed_usernames)}") return logs, len(failed_usernames) > 0