diff --git a/pyproject.toml b/pyproject.toml index 4aa512b..91326be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "ssb-konjunk" -version = "0.1.16" +version = "0.1.17" description = "SSB Konjunk" authors = ["Edvard Garmannslund "] license = "MIT" diff --git a/src/ssb_konjunk/prompts.py b/src/ssb_konjunk/prompts.py index 1bfaa7e..c50624b 100644 --- a/src/ssb_konjunk/prompts.py +++ b/src/ssb_konjunk/prompts.py @@ -271,6 +271,59 @@ def validate_day(day: int | str) -> str: day = "0" + str(int(day)) return str(day) + +def quarter_for_month(month: str | int) -> int: + """Find corresponding quarter for a month. + + Args: + month: Month to find corresponding quarter for. + + Returns: + int: The corresponding quarter. + + Raises: + ValueError: If invalid month + """ + month = int(month) + + if month < 1 or month > 12: + raise ValueError(f"Invalid month: {month}") + + if month < 4: + return 1 + elif month < 7: + return 2 + elif month < 10: + return 3 + else: + return 4 + + +def months_in_quarter(quarter: int | str) -> list[int]: + """Return the three months in the quarter. + + Args: + quarter: the relevant quarter. + + Returns: + list: a list with the months in the quarter. + + Raises: + ValueError: If invalid quarter. + """ + quarter = int(quarter) + + if quarter < 1 or quarter > 4: + raise ValueError(f"Invalid quarter: {quarter}") + + if quarter == 1: + return [1, 2, 3] + elif quarter == 2: + return [4, 5, 6] + elif quarter == 3: + return [7, 8, 9] + else: + return [10, 11, 12] def set_publishing_date() -> str: """Set the date for publication of tables. @@ -355,4 +408,4 @@ def get_previous_month(year: str | int, month: str | int) -> list[int]: prev_month = 12 prev_year = int(year) - 1 - return [prev_year, prev_month] + return [prev_year, prev_month] \ No newline at end of file diff --git a/tests/test_prompts.py b/tests/test_prompts.py index 9ba2f42..5f814c5 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -7,6 +7,7 @@ from ssb_konjunk.prompts import extract_start_end_dates from ssb_konjunk.prompts import get_previous_month from ssb_konjunk.prompts import iterate_years_months +from ssb_konjunk.prompts import quarter_for_month from ssb_konjunk.prompts import validate_month """Test of function days in month""" @@ -151,6 +152,17 @@ def test_validate_month() -> None: assert validate_month("10") == "10" +def test_quarter_for_month() -> None: + # Test with valid integer month + assert quarter_for_month(2) == 1 + + # Test with valid string month + assert quarter_for_month("12") == 4 + + # Test with invalid month + with pytest.raises(ValueError): + quarter_for_month(13) + def test_get_previous_month() -> None: prev_month = get_previous_month(2022, 1) assert prev_month[0] == 2021, f"Previous year for previous month: {prev_month[0]}"