Merge pull request #79 from statisticsnorway/quarter_func

Quarter func

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 <[email protected]>"]
 license = "MIT"

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]

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]}"