move the columns option out of the table formatter and into the lister base

2 files changed, 35 insertions, 14 deletions

diff --git a/cliff/formatters/table.py b/cliff/formatters/table.py
index 690ae30..056557f 100644
--- a/cliff/formatters/table.py
+++ b/cliff/formatters/table.py
@@ -19,18 +19,7 @@ class TableFormatter(ListFormatter, SingleFormatter):
         pass
 
     def add_argument_group(self, parser):
-        group = parser.add_argument_group(
-            title='table formatter',
-            description='Pretty-print output in a table',
-            )
-        group.add_argument(
-            '-c', '--column',
-            action='append',
-            default=[],
-            dest='columns',
-            metavar='COLUMN',
-            help='specify the column(s) to include, can be repeated',
-            )
+        pass
 
     def emit_list(self, column_names, data, stdout, parsed_args):
         x = prettytable.PrettyTable(column_names)
@@ -51,7 +40,7 @@ class TableFormatter(ListFormatter, SingleFormatter):
             x.add_row(first_row)
             for row in data_iter:
                 x.add_row(row)
-        formatted = x.get_string(fields=(parsed_args.columns or column_names))
+        formatted = x.get_string(fields=column_names)
         stdout.write(formatted)
         stdout.write('\n')
         return
diff --git a/cliff/lister.py b/cliff/lister.py
index a4ec7c3..24e3127 100644
--- a/cliff/lister.py
+++ b/cliff/lister.py
@@ -1,6 +1,7 @@
 """Application base class for providing a list of data as output.
 """
 import abc
+import itertools
 import logging
 
 from .display import DisplayCommandBase
@@ -28,8 +29,39 @@ class Lister(DisplayCommandBase):
         containing the data to be listed.
         """
 
+    def get_parser(self, prog_name):
+        parser = super(Lister, self).get_parser(prog_name)
+        parser.add_argument(
+            '-c', '--column',
+            action='append',
+            default=[],
+            dest='columns',
+            metavar='COLUMN',
+            help='specify the column(s) to include, can be repeated',
+            )
+        return parser
+
     def run(self, parsed_args):
         column_names, data = self.get_data(parsed_args)
+        if not parsed_args.columns:
+            columns_to_include = column_names
+            data_gen = data
+        else:
+            columns_to_include = [c for c in column_names
+                                  if c in parsed_args.columns
+                                  ]
+            if not columns_to_include:
+                raise ValueError('No recognized column names in %s' %
+                                 str(parsed_args.columns))
+            # Set up argument to compress()
+            selector = [(c in columns_to_include)
+                        for c in column_names]
+            # Generator expression to only return the parts of a row
+            # of data that the user has expressed interest in
+            # seeing. We have to convert the compress() output to a
+            # list so the table formatter can ask for its length.
+            data_gen = (list(itertools.compress(row, selector))
+                        for row in data)
         formatter = self.formatters[parsed_args.formatter]
-        formatter.emit_list(column_names, data, self.app.stdout, parsed_args)
+        formatter.emit_list(columns_to_include, data_gen, self.app.stdout, parsed_args)
         return 0