From c4ecf8c56fdf980e00bfb4e3d7dc6562266c45e8 Mon Sep 17 00:00:00 2001 From: Kiran Patel <7103956+kiran94@users.noreply.github.com> Date: Fri, 5 May 2023 19:29:06 +0100 Subject: [PATCH] feat(output): allow output file override (#4) * chore(ignore): add csv and parquet files to ignore * chore(ci): add poetry check * feat(output): allow output file override --- .github/workflows/main.yml | 3 +++ .gitignore | 3 +++ prfiesta/__main__.py | 6 ++++-- prfiesta/output.py | 19 ++++++++++--------- tests/test_main.py | 2 +- tests/test_output.py | 16 +++++----------- 6 files changed, 26 insertions(+), 23 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index feb9dd4..30e0b34 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -26,6 +26,9 @@ jobs: with: poetry-version: ${{ matrix.poetry-version }} + - name: Poetry Check + run: poetry check + - name: Install Dependencies run: | make export_requirements diff --git a/.gitignore b/.gitignore index 0f5ea77..c92aacb 100644 --- a/.gitignore +++ b/.gitignore @@ -219,3 +219,6 @@ terraform.rc requirements.txt requirements-dev.txt + +*.csv +*.parquet diff --git a/prfiesta/__main__.py b/prfiesta/__main__.py index c024210..3c919cc 100644 --- a/prfiesta/__main__.py +++ b/prfiesta/__main__.py @@ -20,7 +20,8 @@ @click.option('-u', '--users', required=True, multiple=True, help='The GitHub Users to search for. Can be multiple (space delimited)') @click.option('-t', '--token', help='The Authentication token to use') @click.option('-x', '--url', help='The URL of the Git provider to use') -@click.option('-o', '--output_type', type=click.Choice(['csv', 'parquet']), default='csv', help='The output format') +@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('--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: @@ -28,6 +29,7 @@ def main(**kwargs) -> None: users: tuple[str] = kwargs.get('users') token: str = kwargs.get('token') or github_environment.get_token() url: str = kwargs.get('url') or github_environment.get_url() + output: str = kwargs.get('output') output_type: str = kwargs.get('output_type') before: datetime = kwargs.get('before') after: datetime = kwargs.get('after') @@ -44,7 +46,7 @@ def main(**kwargs) -> None: logger.info('Found [bold green]%s[/bold green] pull requests!', pr_frame.shape[0]) if not pr_frame.empty: - output_frame(pr_frame, output_type, spinner=spinner) + output_frame(pr_frame, output_type, spinner=spinner, output_name=output) if __name__ == '__main__': # pragma: nocover main() diff --git a/prfiesta/output.py b/prfiesta/output.py index dc64469..2c23020 100644 --- a/prfiesta/output.py +++ b/prfiesta/output.py @@ -1,5 +1,4 @@ import logging -import os from datetime import datetime from typing import Literal @@ -10,19 +9,21 @@ logger = logging.getLogger(__name__) -output_directory = 'output' OUTPUT_TYPE = Literal['csv', 'parquet'] -def output_frame(frame: pd.DataFrame, output_type: OUTPUT_TYPE, spinner: Spinner, output_name: str = 'export', timestamp: datetime = None) -> None: +def output_frame( + frame: pd.DataFrame, + output_type: OUTPUT_TYPE, + spinner: Spinner, + output_name: str = None, + timestamp: datetime = None) -> None: - if not timestamp: - timestamp = datetime.now() + if not output_name: + if not timestamp: + timestamp = datetime.now() - os.makedirs(output_directory, exist_ok=True) - - output_name = str(output_name) + '.' + timestamp.strftime('%Y-%m-%d_%H:%M:%S') + '.' + output_type - output_name = os.path.join(output_directory, output_name) + output_name = f"export.{timestamp.strftime('%Y-%m-%d_%H:%M:%S')}.{output_type}" spinner.update(text=f'Writing export to {output_name}', style=SPINNER_STYLE) diff --git a/tests/test_main.py b/tests/test_main.py index b48221f..be2c8d2 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -89,6 +89,6 @@ def test_main( assert mock_collector.return_value.collect.call_args_list == expected_collect_params if not collect_response.empty: - assert mock_output_frame.call_args_list == [call(collect_response, expected_output_type, spinner=mock_spinner.return_value)] + assert mock_output_frame.call_args_list == [call(collect_response, expected_output_type, spinner=mock_spinner.return_value, output_name=None)] assert result.exit_code == 0 diff --git a/tests/test_output.py b/tests/test_output.py index 27df5b5..baf510a 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -1,6 +1,5 @@ -import os from datetime import datetime -from unittest.mock import ANY, Mock, call, patch +from unittest.mock import ANY, Mock, call import pytest @@ -13,23 +12,20 @@ ('csv', None), ('parquet', None), ]) -@patch('prfiesta.output.os') -def test_output_frame(mock_os: Mock, output_type: str, timestamp: datetime) -> None: +def test_output_frame(output_type: str, timestamp: datetime) -> None: mock_frame: Mock = Mock() mock_spinner: Mock = Mock() - mock_os.path.join = os.path.join output_frame(mock_frame, output_type, mock_spinner, timestamp=timestamp) - assert [call('output', exist_ok=True)] == mock_os.makedirs.call_args_list assert mock_spinner.update.called if timestamp and output_type == 'csv': - assert [call('output/export.2021-01-01_00:00:00.csv', index=False)] == mock_frame.to_csv.call_args_list + assert [call('export.2021-01-01_00:00:00.csv', index=False)] == mock_frame.to_csv.call_args_list elif timestamp and output_type == 'parquet': - assert [call('output/export.2021-01-01_00:00:00.parquet', index=False)] == mock_frame.to_parquet.call_args_list + assert [call('export.2021-01-01_00:00:00.parquet', index=False)] == mock_frame.to_parquet.call_args_list elif not timestamp and output_type == 'csv': assert [call(ANY, index=False)] == mock_frame.to_csv.call_args_list @@ -37,12 +33,10 @@ def test_output_frame(mock_os: Mock, output_type: str, timestamp: datetime) -> N elif not timestamp and output_type == 'parquet': assert [call(ANY, index=False)] == mock_frame.to_parquet.call_args_list -@patch('prfiesta.output.os') -def test_output_frame_unknown_type(mock_os: Mock) -> None: +def test_output_frame_unknown_type() -> None: mock_frame: Mock = Mock() mock_spinner: Mock = Mock() - mock_os.path.join = os.path.join with pytest.raises(ValueError, match='unknown output_type'): output_frame(mock_frame, 'unknown_type', mock_spinner, timestamp=datetime(2021, 1, 1))