From 617b8f5ea2703b90fce38bfbe353b76a504ca7bd Mon Sep 17 00:00:00 2001
From: Russell Bryant <rbryant@redhat.com>
Date: Wed, 4 Dec 2013 18:25:13 -0500
Subject: [PATCH] Add --csv-rows option

Add an option to include only --csv-rows <rows> number of rows in the
CSV output.

Change-Id: Ib6d8bf20561b42afef119533b571d7f1f32739a9
---
 reviewstats/cmd/reviewers.py |  7 ++++++-
 tests/test_reviewers.py      | 37 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_reviewers.py

diff --git a/reviewstats/cmd/reviewers.py b/reviewstats/cmd/reviewers.py
index 1710617..863c3dc 100755
--- a/reviewstats/cmd/reviewers.py
+++ b/reviewstats/cmd/reviewers.py
@@ -111,11 +111,13 @@ def write_csv(reviewer_data, file_obj, options, reviewers, projects,
     if ENABLE_RECEIVED:
         row.append('Received')
     writer.writerow(row)
-    for (name, r_data, d_data, s_data) in reviewer_data:
+    for i, (name, r_data, d_data, s_data) in enumerate(reviewer_data, start=1):
         row = [name, r_data, d_data]
         if ENABLE_RECEIVED:
             row.append(s_data)
         writer.writerow(row)
+        if options.csv_rows and i == options.csv_rows:
+            break
 
 
 def write_pretty(reviewer_data, file_obj, options, reviewers, projects,
@@ -241,6 +243,9 @@ def main(argv=None):
         '-u', '--user', default=getpass.getuser(), help='gerrit user')
     optparser.add_option(
         '-k', '--key', default=None, help='ssh key for gerrit')
+    optparser.add_option(
+        '-r', '--csv-rows', default=0, help='Max rows for CSV output',
+        type='int', dest='csv_rows')
 
     options, args = optparser.parse_args()
 
diff --git a/tests/test_reviewers.py b/tests/test_reviewers.py
new file mode 100644
index 0000000..becf6d0
--- /dev/null
+++ b/tests/test_reviewers.py
@@ -0,0 +1,37 @@
+# Copyright (C) 2013 - Red Hat, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import six
+
+import reviewstats.cmd.reviewers
+
+
+class ReviewersCSVTestCase(unittest.TestCase):
+    def test_csv_rows(self):
+        class Options(object):
+            pass
+
+        options = Options()
+        options.csv_rows = 10
+        reviewer_data = [('', '', '', '')] * 100
+        sio = six.StringIO()
+
+        reviewstats.cmd.reviewers.write_csv(reviewer_data, sio, options, {},
+                                            {}, {}, {})
+
+        # NOTE(russellb) With csv_rows set to 10, the output should have 11
+        # lines: a heading line plus 10 rows
+        self.assertEqual(sio.getvalue().count('\n'), 11)