feat(csvclean): Add --fill-short-rows and --fillvalue options

Author: jpmckinney

CHANGELOG.rst

@@ -13,6 +13,8 @@ Other changes:
 
    * :code:`--header-normalize-space`, to strip leading and trailing whitespace and replace sequences of whitespace characters by a single space in the header
    * :code:`--separator`, to change the string with which to join short rows
+   * :code:`--fill-short-rows`, to fill short rows with the missing cells
+   * :code:`--fillvalue`, to change the value with which to fill short rows
 
 * feat: The :code:`--quoting` option accepts 4 (`csv.QUOTE_STRINGS <https://docs.python.org/3/library/csv.html#csv.QUOTE_STRINGS>`__) and 5 (`csv.QUOTE_NOTNULL <https://docs.python.org/3/library/csv.html#csv.QUOTE_NOTNULL>`__) on Python 3.12.
 * feat: :doc:`/scripts/csvformat`: The :code:`--out-quoting` option accepts 4 (`csv.QUOTE_STRINGS <https://docs.python.org/3/library/csv.html#csv.QUOTE_STRINGS>`__) and 5 (`csv.QUOTE_NOTNULL <https://docs.python.org/3/library/csv.html#csv.QUOTE_NOTNULL>`__) on Python 3.12.

csvkit/cleanup.py

@@ -28,10 +28,20 @@ class RowChecker:
     Iterate over rows of a CSV producing cleaned rows and storing error rows.
     """
 
-    def __init__(self, reader, header_normalize_space=False, join_short_rows=False, separator='\n'):
+    def __init__(
+        self,
+        reader,
+        header_normalize_space=False,
+        join_short_rows=False,
+        separator='\n',
+        fill_short_rows=False,
+        fillvalue=None,
+    ):
         self.reader = reader
         self.join_short_rows = join_short_rows
         self.separator = separator
+        self.fill_short_rows = fill_short_rows
+        self.fillvalue = fillvalue
 
         try:
             self.column_names = next(reader)
@@ -50,43 +60,50 @@ def checked_rows(self):
         joinable_row_errors = []
 
         for row in self.reader:
-            if len(row) == length:
+            row_length = len(row)
+
+            if row_length == length:
                 yield row
 
-                # Don't join rows across valid rows.
-                joinable_row_errors = []
+                if self.join_short_rows:
+                    # Don't join rows across valid rows.
+                    joinable_row_errors = []
+
+                continue
+
+            if self.fill_short_rows and row_length < length:
+                yield row + [self.fillvalue] * (length - row_length)
+
                 continue
 
             length_mismatch_error = LengthMismatchError(self.reader.line_num - 1, row, length)
 
             self.errors.append(length_mismatch_error)
 
-            if len(row) > length:
-                # Don't join with long rows.
-                joinable_row_errors = []
-                continue
-
-            if not self.join_short_rows:
-                continue
+            if self.join_short_rows:
+                if row_length > length:
+                    # Don't join with long rows.
+                    joinable_row_errors = []
+                    continue
 
-            joinable_row_errors.append(length_mismatch_error)
-            if len(joinable_row_errors) == 1:
-                continue
+                joinable_row_errors.append(length_mismatch_error)
+                if len(joinable_row_errors) == 1:
+                    continue
 
-            while joinable_row_errors:
-                fixed_row = join_rows([error.row for error in joinable_row_errors], separator=self.separator)
+                while joinable_row_errors:
+                    fixed_row = join_rows([error.row for error in joinable_row_errors], separator=self.separator)
 
-                if len(fixed_row) < length:
-                    break
+                    if len(fixed_row) < length:
+                        break
 
-                if len(fixed_row) == length:
-                    yield fixed_row
+                    if len(fixed_row) == length:
+                        yield fixed_row
 
-                    for fixed in joinable_row_errors:
-                        self.errors.remove(fixed)
+                        for fixed in joinable_row_errors:
+                            self.errors.remove(fixed)
 
-                    joinable_row_errors = []
-                    break
+                        joinable_row_errors = []
+                        break
 
-                # keep trying in case we're too long because of a straggler
-                joinable_row_errors = joinable_row_errors[1:]
+                    # keep trying in case we're too long because of a straggler
+                    joinable_row_errors = joinable_row_errors[1:]

csvkit/utilities/csvclean.py

@@ -23,18 +23,29 @@ def add_arguments(self):
         self.argparser.add_argument(
             '--separator', dest='separator', default='\n',
             help='The string with which to join short rows. Defaults to a newline.')
+        self.argparser.add_argument(
+            '--fill-short-rows', dest='fill_short_rows', action='store_true',
+            help='Fill short rows with the missing cells.')
+        self.argparser.add_argument(
+            '--fillvalue', dest='fillvalue',
+            help='The value with which to fill short rows. Defaults to none.')
 
     def main(self):
         if self.additional_input_expected():
             sys.stderr.write('No input file or piped data provided. Waiting for standard input:\n')
 
+        if self.args.join_short_rows and self.args.fill_short_rows:
+            self.argparser.error('The --join-short-rows and --fill-short-rows options are mutually exclusive.')
+
         reader = agate.csv.reader(self.skip_lines(), **self.reader_kwargs)
 
         checker = RowChecker(
             reader,
             header_normalize_space=self.args.header_normalize_space,
             join_short_rows=self.args.join_short_rows,
             separator=self.args.separator,
+            fill_short_rows=self.args.fill_short_rows,
+            fillvalue=self.args.fillvalue,
         )
 
         output_writer = agate.csv.writer(self.output_file, **self.writer_kwargs)

docs/scripts/csvclean.rst

@@ -36,6 +36,30 @@ Cleans a CSV file of common syntax errors:
       1,"1 Main St, Springfield",US
       2,"123 Acadia Avenue, London",GB
 
+-  If a CSV has missing delimiters, like:
+
+   .. code-block:: none
+
+      id,name,country
+      1,Alice
+      2,Bob,CA
+
+   You can add the missing delimiters with :code:`--fill-short-rows`:
+
+   .. code-block:: none
+
+      id,name,country
+      1,Alice,
+      2,Bob,CA
+
+   To change the value used to fill short rows, use :code:`--fillvalue`. For example, with :code:`--fillvalue "US"`:
+
+   .. code-block:: none
+
+      id,name,country
+      1,Alice,US
+      2,Bob,CA
+
 All valid rows are written to standard output, and all error rows along with line numbers and descriptions are written to standard error. If there are error rows, the exit code will be 1.
 
 .. note::
@@ -63,6 +87,18 @@ All valid rows are written to standard output, and all error rows along with lin
 
    optional arguments:
      -h, --help            show this help message and exit
+     --header-normalize-space
+                           Strip leading and trailing whitespace and replace
+                           sequences of whitespace characters by a single space
+                           in the header.
+     --join-short-rows     Merges short rows into a single row.
+     --separator SEPARATOR
+                           The string with which to join short rows. Defaults to
+                           a newline.
+     --fill-short-rows     Fill short rows with the missing cells.
+     --fillvalue FILLVALUE
+                           The value with which to fill short rows. Defaults to
+                           none.
 
 See also: :doc:`../common_arguments`.
 

examples/test_join_short_rows.csv

@@ -1,4 +1,4 @@
 a,b,c
 1,cat
-dog,3
+dog,c
 3,b,c

tests/test_utilities/test_csvclean.py

@@ -86,14 +86,30 @@ def test_header_normalize_space(self):
     def test_join_short_rows(self):
         self.assertCleaned(['--join-short-rows', 'examples/test_join_short_rows.csv'], [
             ['a', 'b', 'c'],
-            ['1', 'cat\ndog', '3'],
+            ['1', 'cat\ndog', 'c'],
             ['3', 'b', 'c'],
         ])
 
     def test_join_short_rows_separator(self):
         self.assertCleaned(['--join-short-rows', '--separator', 'XYZ', 'examples/test_join_short_rows.csv'], [
             ['a', 'b', 'c'],
-            ['1', 'catXYZdog', '3'],
+            ['1', 'catXYZdog', 'c'],
+            ['3', 'b', 'c'],
+        ])
+
+    def test_fill_short_rows(self):
+        self.assertCleaned(['--fill-short-rows', 'examples/test_join_short_rows.csv'], [
+            ['a', 'b', 'c'],
+            ['1', 'cat', ''],
+            ['dog', 'c', ''],
+            ['3', 'b', 'c'],
+        ])
+
+    def test_fill_short_rows_separator(self):
+        self.assertCleaned(['--fill-short-rows', '--fillvalue', 'XYZ', 'examples/test_join_short_rows.csv'], [
+            ['a', 'b', 'c'],
+            ['1', 'cat', 'XYZ'],
+            ['dog', 'c', 'XYZ'],
             ['3', 'b', 'c'],
         ])