Repository: lzakharov/csv2md

Issue:
Ability to filter columns
It would be great if there was an easy way to provide a list of columns to be used when converting to markdown

So, for example, with a table with 10 columns, doing something like 

    --use-only=1-3,4,5,7


would not include columns 6, 8 9 and 10 into the generated table

Comments:

Patch:
diff --git a/csv2md/__main__.py b/csv2md/__main__.py
index 2f23966..5147845 100644
--- a/csv2md/__main__.py
+++ b/csv2md/__main__.py
@@ -2,16 +2,20 @@ import argparse
 import sys
 
 from .table import Table
+from .exceptions import BaseError, ColumnIdentifierError
 
 
 def main():
-    parser = argparse.ArgumentParser(description='Parse CSV files into Markdown tables.')
+    parser = argparse.ArgumentParser(
+        description='Parse CSV files into Markdown tables.')
     parser.add_argument('files', metavar='CSV_FILE', type=argparse.FileType('r'), nargs='*',
                         help='One or more CSV files to parse')
     parser.add_argument('-d', '--delimiter', metavar='DELIMITER', type=str, default=',',
                         help='delimiter character. Default is \',\'')
     parser.add_argument('-q', '--quotechar', metavar='QUOTECHAR', type=str, default='"',
                         help='quotation character. Default is \'"\'')
+    parser.add_argument('-C', '--columns', dest='columns', type=str, default=None,
+                        help='comma-separated list of column indices or ranges (from zero) to be processed, e.g. "0,3-5,7". Indices out of range will be ignored')
     parser.add_argument('-c', '--center-aligned-columns', metavar='CENTER_ALIGNED_COLUMNS', nargs='*',
                         type=int, default=[], help='column numbers with center alignment (from zero)')
     parser.add_argument('-r', '--right-aligned-columns', metavar='RIGHT_ALIGNED_COLUMNS', nargs='*',
@@ -20,15 +24,39 @@ def main():
                         help='specify that the input CSV file has no header row. Will create default headers in Excel format (a,b,c,...)')
     args = parser.parse_args()
 
-    if not args.files:
-        table = Table.parse_csv(sys.stdin, args.delimiter, args.quotechar)
-        print(table.markdown(args.center_aligned_columns, args.right_aligned_columns, args.no_header_row))
-        return
+    try:
+        columns = parse_columns(args.columns)
+    except BaseError as e:
+        parser.error(e)
 
-    for file in args.files:
-        table = Table.parse_csv(file, args.delimiter, args.quotechar)
+    for file in [sys.stdin] if not args.files else args.files:
+        table = Table.parse_csv(file, args.delimiter, args.quotechar, columns)
         print(table.markdown(args.center_aligned_columns, args.right_aligned_columns, args.no_header_row))
 
 
+def parse_columns(columns):
+    if not columns:
+        return None
+
+    result = []
+
+    for c in columns.split(','):
+        if '-' in c:
+            try:
+                a, b = map(int, c.split('-', 1))
+            except ValueError:
+                raise ColumnIdentifierError(c)
+
+            result.extend(range(a, b + 1))
+        else:
+            if not c.isdecimal():
+                raise ColumnIdentifierError(c)
+
+            column = int(c)
+            result.append(column)
+
+    return result
+
+
 if __name__ == '__main__':
     main()
diff --git a/csv2md/exceptions.py b/csv2md/exceptions.py
new file mode 100644
index 0000000..1fa4de4
--- /dev/null
+++ b/csv2md/exceptions.py
@@ -0,0 +1,8 @@
+class BaseError(Exception):
+    pass
+
+
+class ColumnIdentifierError(BaseError):
+    def __init__(self, column):
+        msg = f'Invalid column identifier "{column}". Must be non-negative integer or range of non-negative integers separated by "-".'
+        super().__init__(msg)
diff --git a/csv2md/table.py b/csv2md/table.py
index 4452cfc..36a9ae1 100644
--- a/csv2md/table.py
+++ b/csv2md/table.py
@@ -37,8 +37,15 @@ class Table:
         return '\n'.join(rows)
 
     @staticmethod
-    def parse_csv(file, delimiter=',', quotechar='"'):
-        return Table(list(csv.reader(file, delimiter=delimiter, quotechar=quotechar)))
+    def parse_csv(file, delimiter=',', quotechar='"', columns=None):
+        reader = csv.reader(file, delimiter=delimiter, quotechar=quotechar)
+
+        if columns is None:
+            cells = list(reader)
+        else:
+            cells = [[row[i] for i in columns if 0 <= i < len(row)] for row in reader]
+
+        return Table(cells)
 
     @staticmethod
     def make_default_headers(n):


Test patch:
diff --git a/csv2md/test_table.py b/csv2md/test_table.py
index a6128cf..0ca823e 100644
--- a/csv2md/test_table.py
+++ b/csv2md/test_table.py
@@ -26,6 +26,15 @@ normal_cells = [
 
 normal_widths = [4, 5, 26, 33, 7]
 
+filtered_columns_cells = [
+    ['year', 'model', 'description'],
+    ['1997', 'E350', 'ac, abs, moon'],
+    ['1999', 'Venture «Extended Edition»', ''],
+    ['1996', 'Grand Cherokee', 'MUST SELL! air, moon roof, loaded']
+]
+
+filtered_columns_widths = [4, 26, 33]
+
 normal_md = (
     '| year | make  | model                      | description                       | price   |\n'
     '| ---- | ----- | -------------------------- | --------------------------------- | ------- |\n'
@@ -91,6 +100,20 @@ class TestTable(TestCase):
         self.assertEqual(expected_cells, actual.cells)
         self.assertEqual(expected_widths, actual.widths)
 
+    def test_parse_csv_with_columns(self):
+        expected_cells = filtered_columns_cells
+        expected_widths = filtered_columns_widths
+        actual = Table.parse_csv(io.StringIO(normal_csv), columns=[0, 2, 3])
+        self.assertEqual(expected_cells, actual.cells)
+        self.assertEqual(expected_widths, actual.widths)
+
+    def test_parse_csv_with_invalid_columns(self):
+        expected_cells = filtered_columns_cells
+        expected_widths = filtered_columns_widths
+        actual = Table.parse_csv(io.StringIO(normal_csv), columns=[-10, -1, 0, 2, 3, 5, 10, 100])
+        self.assertEqual(expected_cells, actual.cells)
+        self.assertEqual(expected_widths, actual.widths)
+
     def test_make_default_headers(self):
         expected = ('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l',
                     'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x',