Add support for Excel sheet name regexps

Author: rachekalmir

csvkit/utilities/in2csv.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-
+import re
 import sys
 from os.path import splitext
 
@@ -47,7 +47,9 @@ def option_parser(bytestring):
         self.argparser.add_argument('--write-sheets', dest='write_sheets', type=option_parser,
                                     help='The names of the Excel sheets to write to files, or "-" to write all sheets.')
         self.argparser.add_argument('--write-sheet-names', dest='write_sheet_names', action='store_true',
-                                    help='Use the sheet name to create the target filename.')
+                                    help='Use the sheet name to create the target filename when using --write-sheets.')
+        self.argparser.add_argument('-r', '--regular-expression', dest='regular_expression',
+                                    help='Use regular expression to identify XLS sheet names when using --write-sheets.')
         self.argparser.add_argument('--encoding-xls', dest='encoding_xls',
                                     help='Specify the encoding of the input XLS file.')
         self.argparser.add_argument('-y', '--snifflimit', dest='sniff_limit', type=int,
@@ -162,6 +164,8 @@ def main(self):
 
             if self.args.write_sheets == '-':
                 sheets = self.sheet_names(path, filetype)
+            elif self.args.regular_expression:
+                sheets = [sheet for sheet in self.sheet_names(path, filetype) if re.match(self.args.write_sheets, sheet)]
             else:
                 sheets = [int(sheet) if sheet.isdigit() else sheet for sheet in self.args.write_sheets.split(',')]
 

docs/scripts/in2csv.rst

@@ -47,7 +47,9 @@ The header line is required though the columns may be in any order::
       --write-sheets WRITE_SHEETS
                             The names of the Excel sheets to write to files, or
                             "-" to write all sheets.
-      --write-sheet-names   Use the sheet name to create the target filename.
+      --write-sheet-names   Use the sheet name to create the target filename when using --write-sheets.
+      -r, --regular-expression
+                            Use regular expression to identify Excel sheet names when using --write-sheets.
       -y SNIFF_LIMIT, --snifflimit SNIFF_LIMIT
                             Limit CSV dialect sniffing to the specified number of
                             bytes. Specify "0" to disable sniffing entirely.

tests/test_utilities/test_in2csv.py

@@ -261,3 +261,18 @@ def test_convert_xlsx_with_write_sheets_with_names(self):
                 path = 'examples/sheets_%s.csv' % suffix
                 if os.path.exists(path):
                     os.remove(path)
+
+    def test_convert_xlsx_with_write_sheets_with_names_and_regexp(self):
+        try:
+            self.assertConverted('xlsx', 'examples/sheets.xlsx', 'examples/testxlsx_noinference_converted.csv', ['--no-inference', '-r', '--write-sheets', "[a-z]+", '--write-sheet-names'])
+            with open('examples/sheets_data.csv', 'r') as f:
+                with open('examples/testxlsx_noinference_converted.csv', 'r') as g:
+                    self.assertEqual(f.read(), g.read())
+            self.assertFalse(os.path.exists('examples/sheets_0.csv'))
+            self.assertFalse(os.path.exists('examples/sheets_1.csv'))
+            self.assertFalse(os.path.exists('examples/sheets_2.csv'))
+        finally:
+            for suffix in ('data',):
+                path = 'examples/sheets_%s.csv' % suffix
+                if os.path.exists(path):
+                    os.remove(path)