171 lines
3.8 KiB
Python
171 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import csv
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def split_cell(value):
|
|
"""
|
|
Split a cell on slash characters.
|
|
|
|
Empty parts are removed, so:
|
|
"A / B" becomes ["A", "B"]
|
|
"""
|
|
|
|
if value is None:
|
|
return [""]
|
|
|
|
parts = re.split(r"\s*/\s*", value.strip())
|
|
|
|
return [
|
|
part.strip()
|
|
for part in parts
|
|
if part.strip()
|
|
]
|
|
|
|
|
|
def expand_row(row, line_number):
|
|
"""
|
|
Expand one CSV row into one or more rows.
|
|
|
|
Cells without a slash are repeated.
|
|
Cells with slashes must contain the same number
|
|
of values as the other split cells.
|
|
"""
|
|
|
|
split_values = {
|
|
column: split_cell(value)
|
|
for column, value in row.items()
|
|
}
|
|
|
|
split_lengths = [
|
|
len(values)
|
|
for values in split_values.values()
|
|
if len(values) > 1
|
|
]
|
|
|
|
# A row with no slash produces one row.
|
|
if not split_lengths:
|
|
return [row]
|
|
|
|
expected_count = split_lengths[0]
|
|
|
|
# Every cell containing a slash must split into the same
|
|
# number of values.
|
|
invalid_columns = [
|
|
column
|
|
for column, values in split_values.items()
|
|
if len(values) not in (1, expected_count)
|
|
]
|
|
|
|
if invalid_columns:
|
|
raise ValueError(
|
|
f"line {line_number}: inconsistent number of values. "
|
|
f"Expected {expected_count} values in slash-separated "
|
|
f"cells, but these columns differ: "
|
|
f"{', '.join(invalid_columns)}"
|
|
)
|
|
|
|
expanded_rows = []
|
|
|
|
for index in range(expected_count):
|
|
expanded_row = {}
|
|
|
|
for column, values in split_values.items():
|
|
if len(values) == 1:
|
|
# Repeat values that did not contain a slash.
|
|
expanded_row[column] = values[0]
|
|
else:
|
|
# Take the corresponding split value.
|
|
expanded_row[column] = values[index]
|
|
|
|
expanded_rows.append(expanded_row)
|
|
|
|
return expanded_rows
|
|
|
|
|
|
def preprocess_csv(input_path, output_path):
|
|
output_rows = []
|
|
|
|
with open(
|
|
input_path,
|
|
"r",
|
|
encoding="utf-8-sig",
|
|
newline="",
|
|
) as input_file:
|
|
reader = csv.DictReader(input_file)
|
|
|
|
if not reader.fieldnames:
|
|
raise ValueError("The input CSV has no header row.")
|
|
|
|
fieldnames = reader.fieldnames
|
|
|
|
for line_number, row in enumerate(reader, start=2):
|
|
# Ignore completely empty rows.
|
|
if not any(
|
|
value is not None and value.strip()
|
|
for value in row.values()
|
|
):
|
|
continue
|
|
|
|
expanded_rows = expand_row(row, line_number)
|
|
output_rows.extend(expanded_rows)
|
|
|
|
with open(
|
|
output_path,
|
|
"w",
|
|
encoding="utf-8",
|
|
newline="",
|
|
) as output_file:
|
|
writer = csv.DictWriter(
|
|
output_file,
|
|
fieldnames=fieldnames,
|
|
extrasaction="ignore",
|
|
)
|
|
|
|
writer.writeheader()
|
|
writer.writerows(output_rows)
|
|
|
|
print(
|
|
f"Created {output_path} with "
|
|
f"{len(output_rows)} row(s).",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Expand slash-separated values in any CSV column "
|
|
"into separate rows."
|
|
)
|
|
)
|
|
|
|
parser.add_argument(
|
|
"input_csv",
|
|
help="Input CSV file",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"output_csv",
|
|
help="Output CSV file",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
preprocess_csv(
|
|
Path(args.input_csv),
|
|
Path(args.output_csv),
|
|
)
|
|
except (OSError, ValueError) as error:
|
|
print(f"Error: {error}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|