feat: enforce data frame dtypes

closes #5

Author: sgratzl

delphi_epidata/_model.py

@@ -2,11 +2,25 @@
 from enum import Enum
 from datetime import date
 from urllib.parse import urlencode
-from typing import Final, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, TypedDict, Union
-
-from pandas import DataFrame, to_datetime
-
-from ._parse import parse_api_date, parse_api_week
+from typing import (
+    Any,
+    Dict,
+    Final,
+    Generic,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    TypedDict,
+    Union,
+)
+
+from pandas import DataFrame, CategoricalDtype
+
+from ._parse import parse_api_date, parse_api_week, fields_to_predicate
 
 EpiRangeDict = TypedDict("EpiRangeDict", {"from": int, "to": int})
 EpiRangeLike = Union[int, str, "EpiRange", EpiRangeDict, date]
@@ -224,17 +238,29 @@ def _parse_row(
     def _as_df(
         self,
         rows: Sequence[Mapping[str, Union[str, float, int, date, None]]],
+        fields: Optional[Iterable[str]] = None,
         disable_date_parsing: Optional[bool] = False,
     ) -> DataFrame:
-        # TODO define data frame dtypes for each column
-        df = DataFrame(rows)
+        pred = fields_to_predicate(fields)
+        columns: List[str] = [info.name for info in self.meta if pred(info.name)]
+        df = DataFrame(rows, columns=columns or None)
+
+        data_types: Dict[str, Any] = {}
         for info in self.meta:
-            if (
-                info.type in (EpidataFieldType.date, EpidataFieldType.epiweek)
-                and info.name in df.columns
-                and not disable_date_parsing
-            ):
-                df[info.name] = to_datetime(df[info.name])
-            if info.type == EpidataFieldType.categorical and info.categories and info.name in df.columns:
-                df[info.name] = df[info.name].astype("category").cat.set_categories(info.categories, ordered=True)
+            if not pred(info.name):
+                continue
+            if info.type == EpidataFieldType.bool:
+                data_types[info.name] = bool
+            elif info.type == EpidataFieldType.categorical:
+                data_types[info.name] = CategoricalDtype(categories=info.categories or None, ordered=True)
+            elif info.type == EpidataFieldType.int:
+                data_types[info.name] = int
+            elif info.type in (EpidataFieldType.date, EpidataFieldType.epiweek):
+                data_types[info.name] = int if disable_date_parsing else "datetime64"
+            elif info.type == EpidataFieldType.float:
+                data_types[info.name] = float
+            else:
+                data_types[info.name] = str
+        if data_types:
+            df = df.astype(data_types)
         return df

delphi_epidata/_parse.py

@@ -1,4 +1,4 @@
-from typing import Optional, cast
+from typing import Callable, Iterable, Optional, Set, cast
 
 from typing import Union
 from datetime import date, datetime
@@ -16,3 +16,16 @@ def parse_api_week(value: Union[str, int, float, None]) -> Optional[date]:
     if value is None:
         return None
     return cast(date, Week.fromstring(str(value)).startdate())
+
+
+def fields_to_predicate(fields: Optional[Iterable[str]] = None) -> Callable[[str], bool]:
+    if not fields:
+        return lambda _: True
+    to_include: Set[str] = set()
+    to_exclude: Set[str] = set()
+    for f in fields:
+        if f.startswith("-"):
+            to_exclude.add(f[1:])
+        else:
+            to_include.add(f)
+    return lambda f: (f not in to_exclude and (not to_include or f in to_include))

delphi_epidata/async_requests.py

@@ -121,7 +121,7 @@ async def df(
         if self.only_supports_classic:
             raise OnlySupportsClassicFormatException()
         r = await self.json(fields, disable_date_parsing=disable_date_parsing)
-        return self._as_df(r, disable_date_parsing)
+        return self._as_df(r, fields, disable_date_parsing)
 
     async def csv(self, fields: Optional[Iterable[str]] = None) -> str:
         """Request and parse epidata in CSV format"""

delphi_epidata/requests.py

@@ -109,8 +109,8 @@ def df(self, fields: Optional[Iterable[str]] = None, disable_date_parsing: Optio
         """Request and parse epidata as a pandas data frame"""
         if self.only_supports_classic:
             raise OnlySupportsClassicFormatException()
-        r = self.json(fields)
-        return self._as_df(r, disable_date_parsing=disable_date_parsing)
+        r = self.json(fields, disable_date_parsing=disable_date_parsing)
+        return self._as_df(r, fields, disable_date_parsing=disable_date_parsing)
 
     def csv(self, fields: Optional[Iterable[str]] = None) -> str:
         """Request and parse epidata in CSV format"""

smoke_test.py

@@ -17,6 +17,11 @@
 df = apicall.df()
 print(df.columns)
 print(df.dtypes)
+print(df.iloc[0])
+df = apicall.df(disable_date_parsing=True)
+print(df.columns)
+print(df.dtypes)
+print(df.iloc[0])
 
 for row in apicall.iter():
     print(row)