refactor: async is gone, consolidate df code

Author: dshemetov

epidatpy/_model.py

@@ -2,8 +2,6 @@
 from datetime import date
 from enum import Enum
 from typing import (
-    Any,
-    Dict,
     Final,
     List,
     Mapping,
@@ -18,10 +16,8 @@
 from urllib.parse import urlencode
 
 from epiweeks import Week
-from pandas import CategoricalDtype, DataFrame, Series
 
 from ._parse import (
-    fields_to_predicate,
     parse_api_date,
     parse_api_date_or_week,
     parse_api_week,
@@ -257,39 +253,3 @@ def _parse_row(
         if not self.meta:
             return row
         return {k: self._parse_value(k, v, disable_date_parsing) for k, v in row.items()}
-
-    def _as_df(
-        self,
-        rows: Sequence[Mapping[str, Union[str, float, int, date, None]]],
-        fields: Optional[Sequence[str]] = None,
-        disable_date_parsing: Optional[bool] = False,
-    ) -> DataFrame:
-        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 not pred(info.name) or df[info.name].isnull().all():
-                continue
-            if info.type == EpidataFieldType.bool:
-                data_types[info.name] = bool
-            elif info.type == EpidataFieldType.categorical:
-                data_types[info.name] = CategoricalDtype(
-                    categories=Series(info.categories) if info.categories else None, ordered=True
-                )
-            elif info.type == EpidataFieldType.int:
-                data_types[info.name] = int
-            elif info.type in (
-                EpidataFieldType.date,
-                EpidataFieldType.epiweek,
-                EpidataFieldType.date_or_epiweek,
-            ):
-                data_types[info.name] = int if disable_date_parsing else "datetime64[ns]"
-            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

epidatpy/request.py

@@ -1,5 +1,7 @@
 from datetime import date
 from typing import (
+    Any,
+    Dict,
     Final,
     List,
     Mapping,
@@ -9,7 +11,7 @@
     cast,
 )
 
-from pandas import DataFrame
+from pandas import CategoricalDtype, DataFrame, Series
 from requests import Response, Session
 from requests.auth import HTTPBasicAuth
 from tenacity import retry, stop_after_attempt
@@ -21,13 +23,15 @@
 from ._model import (
     AEpiDataCall,
     EpidataFieldInfo,
+    EpidataFieldType,
     EpiDataFormatType,
     EpiDataResponse,
     EpiRange,
     EpiRangeParam,
     OnlySupportsClassicFormatException,
     add_endpoint_to_url,
 )
+from ._parse import fields_to_predicate
 
 # Make the linter happy about the unused variables
 __all__ = ["Epidata", "EpiDataCall", "EpiDataContext", "EpiRange", "CovidcastEpidata"]
@@ -140,8 +144,36 @@ def df(
         if self.only_supports_classic:
             raise OnlySupportsClassicFormatException()
         self._verify_parameters()
-        r = self.json(fields, disable_date_parsing=disable_date_parsing)
-        return self._as_df(r, fields, disable_date_parsing=disable_date_parsing)
+        rows = self.json(fields, disable_date_parsing=disable_date_parsing)
+        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 not pred(info.name) or df[info.name].isnull().all():
+                continue
+            if info.type == EpidataFieldType.bool:
+                data_types[info.name] = bool
+            elif info.type == EpidataFieldType.categorical:
+                data_types[info.name] = CategoricalDtype(
+                    categories=Series(info.categories) if info.categories else None, ordered=True
+                )
+            elif info.type == EpidataFieldType.int:
+                data_types[info.name] = "Int64"
+            elif info.type in (
+                EpidataFieldType.date,
+                EpidataFieldType.epiweek,
+                EpidataFieldType.date_or_epiweek,
+            ):
+                data_types[info.name] = "Int64" if disable_date_parsing else "datetime64[ns]"
+            elif info.type == EpidataFieldType.float:
+                data_types[info.name] = "Float64"
+            else:
+                data_types[info.name] = "string"
+        if data_types:
+            df = df.astype(data_types)
+        return df
 
 
 class EpiDataContext(AEpiDataEndpoints[EpiDataCall]):