feat(output): allow custom drop columns (#27)

* feat(collector): allow user to specify drop columns

* refactor(collector): ensure move to end process gracefully handles errors

* refactor(main): wire up drop columns into cli interface

* docs(readme): add drop column example

* test(github): ensure dummy token is passed

Author: kiran94

README.md

@@ -42,6 +42,9 @@ prfiesta -u kiran94 --output_type parquet --output my_pull_requests.parquet
 # Get all pull requests for more then one user
 prfiesta -u kiran94 -u user2
 
+# Get all pull requests and drop specific columns from the output
+prfiesta -u kiran94 -dc events_url -dc comments_url -dc node_id
+
 # Get help
 prfiesta --help
 ```

prfiesta/__main__.py

@@ -22,6 +22,7 @@
 @click.option('-x', '--url', help='The URL of the Git provider to use')
 @click.option('-o', '--output', default=None, help='The output location')
 @click.option('-ot', '--output_type', type=click.Choice(['csv', 'parquet']), default='csv', help='The output format')
+@click.option('-dc', '--drop_columns', multiple=True, help='Drop columns from the output dataframe')
 @click.option('--after', type=click.DateTime(formats=['%Y-%m-%d']), help='Only search for pull requests after this date e.g 2023-01-01')
 @click.option('--before', type=click.DateTime(formats=['%Y-%m-%d']), help='Only search for pull requests before this date e.g 2023-04-30')
 def main(**kwargs) -> None:
@@ -33,14 +34,15 @@ def main(**kwargs) -> None:
     output_type: str = kwargs.get('output_type')
     before: datetime = kwargs.get('before')
     after: datetime = kwargs.get('after')
+    drop_columns: list[str] = list(kwargs.get('drop_columns'))
 
     logger.info('[bold green]PR Fiesta 🦜🥳')
 
     spinner = Spinner('dots', text=Text('Loading', style=SPINNER_STYLE))
 
     with Live(spinner, refresh_per_second=20, transient=True):
 
-        collector = GitHubCollector(token=token, url=url, spinner=spinner)
+        collector = GitHubCollector(token=token, url=url, spinner=spinner, drop_columns=drop_columns)
         pr_frame = collector.collect(*users, after=after, before=before)
 
         if not pr_frame.empty:

prfiesta/collectors/github.py

@@ -25,10 +25,8 @@ def __init__(self, **kwargs) -> None:
         self._spinner: Spinner = kwargs.get('spinner')
 
         self._sort_column = ['updated_at']
-        self._drop_columns = [
-            'node_id',
-            'performed_via_github_app',
-        ]
+        self._drop_columns = kwargs.get('drop_columns') or ['node_id', 'performed_via_github_app']
+
         self._move_to_end_columns = [
             'url',
             'repository_url',
@@ -113,8 +111,13 @@ def _construct_query(users: List[str], after: Optional[datetime] = None, before:
 
     def _move_column_to_end(self, df: pd.DataFrame) -> pd.DataFrame:
         for col in self._move_to_end_columns:
-            df.insert(len(df.columns)-1, col, df.pop(col))
-            df.drop(columns=col)
+            try:
+                df.insert(len(df.columns)-1, col, df.pop(col))
+                df.drop(columns=col)
+            except KeyError:
+                # This can happen if the user provides a custom _drop_columns which
+                # removes the column before we can move it to the end
+                logger.debug('Attempted to move column %s but it did not exist', col)
 
         return df
 

tests/collectors/test_github.py

@@ -159,3 +159,26 @@ def test_collect_rate_limit(mock_github: Mock) -> None:
     result = gc.collect('user')
 
     assert result.empty
+
+
+@patch('prfiesta.collectors.github.Github')
+def test_collect_custom_drop_columns(mock_github: Mock) -> None:
+
+    mock_github.return_value.search_issues.return_value = [_mock_issue1]
+
+    collector_params = {
+        'token': 'dummy_token',
+        'url': 'dummy_url',
+        'drop_columns': ['comments_url'],
+    }
+
+    gc = GitHubCollector(**collector_params)
+    result = gc.collect('user1')
+
+    columns = result.columns.tolist()
+    assert 'comments_url' not in columns
+
+    # These are default drop columns
+    # Since we are overriding it in this scenario, they should still exist in the output column
+    assert 'node_id' in columns
+    assert 'performed_via_github_app' in columns

tests/test_main.py

@@ -99,7 +99,7 @@ def test_main(
     assert mock_live.called
     assert mock_spinner.called
 
-    assert mock_collector.call_args_list == [call(token=ANY, url='https://api.github.com', spinner=mock_spinner.return_value)]
+    assert mock_collector.call_args_list == [call(token=ANY, url='https://api.github.com', spinner=mock_spinner.return_value, drop_columns=[])]
     assert mock_collector.return_value.collect.call_args_list == expected_collect_params
 
     if not collect_response.empty: