Factor out function pick_prices_from_dataframe

That function can be useful by itself for consumers of the package.
ymyke · Apr 6, 2022 · f6ced7e · f6ced7e
1 parent 61d3361
commit f6ced7e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/pypme/__init__.py b/pypme/__init__.py
@@ -1,3 +1,7 @@
 __version__ = '0.2.0'
 from .pme import verbose_pme, pme, verbose_xpme, xpme
-from .mod_investpy_pme import investpy_verbose_pme, investpy_pme
+from .mod_investpy_pme import (
+    investpy_verbose_pme,
+    investpy_pme,
+    pick_prices_from_dataframe,
+)
diff --git a/pypme/mod_investpy_pme.py b/pypme/mod_investpy_pme.py
@@ -21,6 +21,18 @@
 from .pme import verbose_xpme
 
 
+def pick_prices_from_dataframe(
+    dates: List[date], pricedf: pd.DataFrame, which_column: str
+) -> List[float]:
+    """Return the prices from `pricedf` that are nearest to dates `dates`. Use
+    `which_column` to pick the dataframe column.
+    """
+    return list(
+        pricedf.iloc[pricedf.index.get_indexer([x], method="nearest")[0]][which_column]
+        for x in dates
+    )
+
+
 def get_historical_data(ticker: str, type: str, **kwargs) -> pd.DataFrame:
     """Small wrapper to make the investpy interface accessible in a more unified fashion."""
     kwargs[type] = ticker
@@ -47,12 +59,9 @@ def investpy_verbose_pme(
         from_date=dates[0].strftime("%d/%m/%Y"),
         to_date=dates[-1].strftime("%d/%m/%Y"),
     )
-    # Pick the nearest price if there is no price for an exact date:
-    pme_prices = [
-        pmedf.iloc[pmedf.index.get_indexer([x], method="nearest")[0]]["Close"]
-        for x in dates
-    ]
-    return verbose_xpme(dates, cashflows, prices, pme_prices)
+    return verbose_xpme(
+        dates, cashflows, prices, pick_prices_from_dataframe(dates, pmedf, "Close")
+    )
 
 
 def investpy_pme(