diff --git a/csvkit/utilities/csvcut.py b/csvkit/utilities/csvcut.py index 10d42d148..f108a7b54 100644 --- a/csvkit/utilities/csvcut.py +++ b/csvkit/utilities/csvcut.py @@ -35,6 +35,9 @@ def add_arguments(self): self.argparser.add_argument( '-x', '--delete-empty-rows', dest='delete_empty', action='store_true', help='After cutting, delete rows which are completely empty.') + self.argparser.add_argument( + '-N', '--column-numbers', dest='number_columns_everywhere', action='store_true', + help='Add column numbering everywhere, in header and cells.') def main(self): if self.args.names_only: @@ -47,10 +50,19 @@ def main(self): rows, column_names, column_ids = self.get_rows_and_column_names_and_column_ids(**self.reader_kwargs) output = agate.csv.writer(self.output_file, **self.writer_kwargs) - output.writerow([column_names[column_id] for column_id in column_ids]) + + one_or_zero = 0 if self.args.zero_based else 1 + if self.args.number_columns_everywhere: + output.writerow([f'{column_id + one_or_zero}~{column_names[column_id]}' for column_id in column_ids]) + else: + output.writerow([column_names[column_id] for column_id in column_ids]) for row in rows: - out_row = [row[column_id] if column_id < len(row) else None for column_id in column_ids] + if self.args.number_columns_everywhere: + out_row = [f'{column_id + one_or_zero}~{row[column_id]}' + if column_id < len(row) else None for column_id in column_ids] + else: + out_row = [row[column_id] if column_id < len(row) else None for column_id in column_ids] if not self.args.delete_empty or any(out_row): output.writerow(out_row) diff --git a/tests/test_utilities/test_csvcut.py b/tests/test_utilities/test_csvcut.py index c78e98248..e250ae0ef 100644 --- a/tests/test_utilities/test_csvcut.py +++ b/tests/test_utilities/test_csvcut.py @@ -24,6 +24,12 @@ def test_simple(self): ['1', '3'], ]) + def test_column_numbers(self): + self.assertRows(['-c', '1,3', '-N', 'examples/dummy.csv'], [ + ['1~a', '3~c'], + ['1~1', '3~3'], + ]) + def test_linenumbers(self): self.assertRows(['-c', '1,3', '-l', 'examples/dummy.csv'], [ ['line_number', 'a', 'c'],