refactor: EpiRange validates input

Author: dshemetov

epidatpy/_model.py

@@ -3,7 +3,6 @@
 from enum import Enum
 from typing import (
     Final,
-    Generic,
     List,
     Mapping,
     Optional,
@@ -22,6 +21,7 @@
     parse_api_date,
     parse_api_date_or_week,
     parse_api_week,
+    parse_user_date_or_week,
 )
 
 EpiDateLike = Union[int, str, date, Week]
@@ -30,9 +30,8 @@
 EpiRangeParam = Union[EpiRangeLike, Sequence[EpiRangeLike]]
 StringParam = Union[str, Sequence[str]]
 IntParam = Union[int, Sequence[int]]
-EpiDataResponse = TypedDict("EpiDataResponse", {"result": int, "message": str, "epidata": List})
 ParamType = Union[StringParam, IntParam, EpiRangeParam]
-EPI_RANGE_TYPE = TypeVar("EPI_RANGE_TYPE", int, date, str, Week)
+EpiDataResponse = TypedDict("EpiDataResponse", {"result": int, "message": str, "epidata": List})
 CALL_TYPE = TypeVar("CALL_TYPE")
 
 
@@ -66,20 +65,15 @@ def format_list(values: EpiRangeParam) -> str:
     return format_item(values)
 
 
-def format_epiweek(value: Union[str, int]) -> str:
-    return Week.fromstring(str(value)).cdcformat()
-
-
-@dataclass(repr=False)
-class EpiRange(Generic[EPI_RANGE_TYPE]):
+class EpiRange:
     """
     Range object for dates/epiweeks
     """
 
-    start: EPI_RANGE_TYPE
-    end: EPI_RANGE_TYPE
-
-    def __post_init__(self) -> None:
+    def __init__(self, start: EpiDateLike, end: EpiDateLike) -> None:
+        # check if types are correct
+        self.start = parse_user_date_or_week(start)
+        self.end = parse_user_date_or_week(end)
         # swap if wrong order
         # complicated construct for typing inference
         if self.end < self.start:

tests/test_model.py

@@ -1,29 +1,31 @@
+import datetime
+
 from epidatpy._model import EpiRange, format_item, format_list
 
 
 def test_epirange() -> None:
-    r = EpiRange(3, 4)
-    assert r.start == 3 and r.end == 4
-    assert str(r) == "3-4"
+    r = EpiRange(20000101, 20000102)
+    assert r.start == datetime.date(2000, 1, 1) and r.end == datetime.date(2000, 1, 2)
+    assert str(r) == "20000101-20000102"
 
 
 def test_epirange_wrong_order() -> None:
-    r = EpiRange(4, 3)
-    assert r.start == 3 and r.end == 4
+    r = EpiRange(20000101, 20000102)
+    assert r.start == datetime.date(2000, 1, 1) and r.end == datetime.date(2000, 1, 2)
 
 
 def test_format_item() -> None:
     assert format_item("a") == "a"
     assert format_item(1) == "1"
     assert format_item({"from": 1, "to": 3}) == "1-3"
-    assert format_item(EpiRange(3, 5)) == "3-5"
+    assert format_item(EpiRange(20000101, 20000102)) == "20000101-20000102"
 
 
 def test_format_list() -> None:
     assert format_list("a") == "a"
     assert format_list(1) == "1"
     assert format_list({"from": 1, "to": 3}) == "1-3"
-    assert format_list(EpiRange(3, 5)) == "3-5"
+    assert format_list(EpiRange(20000101, 20000102)) == "20000101-20000102"
 
     assert format_list(["a", "b"]) == "a,b"
     assert format_list(("a", "b")) == "a,b"