refactor: convenience functions for CovidcastEpidata

Author: dshemetov

epidatpy/_covidcast.py

@@ -145,7 +145,7 @@ def to_df(signals: Iterable["DataSignal"]) -> DataFrame:
             ],
         )
         df["geo_types"] = [",".join(s.geo_types.keys()) for s in signals]
-        return df.set_index(["source", "signal"])
+        return df
 
     @property
     def key(self) -> Tuple[str, str]:
@@ -236,7 +236,7 @@ def to_df(sources: Iterable["DataSource"]) -> DataFrame:
             ],
         )
         df["signals"] = [",".join(ss.signal for ss in s.signals) for s in sources]
-        return df.set_index("source")
+        return df
 
     def get_signal(self, signal: str) -> Optional[DataSignal]:
         return next((s for s in self.signals if s.signal == signal), None)
@@ -266,10 +266,14 @@ def __post_init__(self) -> None:
             for signal in source.signals:
                 self._signals_by_key[signal.key] = signal
 
-    @property
     def source_names(self) -> Sequence[str]:
         return [s.source for s in self.sources]
 
+    def signal_names(self, source: Optional[str] = None) -> Sequence[str]:
+        if not source:
+            return [x.signal for src in self._source_by_name.values() for x in src.signals]
+        return [s.signal for s in self._source_by_name[source].signals]
+
     @cached_property
     def source_df(self) -> DataFrame:
         """Fetch metadata about available covidcast sources.

smoke_test.py

@@ -1,28 +1,33 @@
 from datetime import date
+from typing import cast
+
+import pandas as pd
 
 from epidatpy import CovidcastEpidata, Epidata, EpiRange
 
 print("Epidata Test")
 apicall = Epidata.pub_covidcast("fb-survey", "smoothed_cli", "nation", "day", "us", EpiRange(20210405, 20210410))
 
+# Call info
 print(apicall)
+# URL
 print(apicall.request_url())
-
-classic = apicall.classic()
-print(classic)
-
+# DataFrame
 df = apicall.df()
 print(df.columns)
 print(df.dtypes)
 print(df.iloc[0])
+print(df)
+# Classic
+classic = apicall.classic()
+# DataFrame
 df = apicall.df(disable_date_parsing=True)
 print(df.columns)
 print(df.dtypes)
 print(df.iloc[0])
 
 
 StagingEpidata = Epidata.with_base_url("https://staging.delphi.cmu.edu/epidata/")
-
 epicall = StagingEpidata.pub_covidcast(
     "fb-survey", "smoothed_cli", "nation", "day", "*", EpiRange(date(2021, 4, 5), date(2021, 4, 10))
 )
@@ -32,23 +37,23 @@
 # Covidcast test
 print("Covidcast Test")
 epidata = CovidcastEpidata()
-print(epidata.source_names)
+epidata.source_df
+epidata.signal_df
+print(epidata.source_names())
+print(epidata.signal_names("hhs"))
+
+# Print signals as df
 epidata["fb-survey"].signal_df
 apicall = epidata[("fb-survey", "smoothed_cli")].call(
     "nation",
     "us",
     EpiRange(20210405, 20210410),
 )
 print(apicall)
-
-classic = apicall.classic()
-print(classic)
-
-df = apicall.df()
-print(df.columns)
-print(df.dtypes)
-print(df.iloc[0])
-df = apicall.df(disable_date_parsing=True)
+df = cast(pd.DataFrame, apicall())
 print(df.columns)
 print(df.dtypes)
 print(df.iloc[0])
+
+classic = apicall.classic()
+print(classic)