From dba627882c70ae5c4c3535f7057073d5ff8734a4 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Thu, 7 Sep 2023 09:24:14 +0200 Subject: [PATCH] fix mypy --- pymc_marketing/clv/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pymc_marketing/clv/utils.py b/pymc_marketing/clv/utils.py index 0eb3ac01..11f12a2d 100644 --- a/pymc_marketing/clv/utils.py +++ b/pymc_marketing/clv/utils.py @@ -168,6 +168,8 @@ def _find_first_transactions( if isinstance(observation_period_end, pd.Period): observation_period_end = observation_period_end.to_timestamp() + if isinstance(observation_period_end, str): + observation_period_end = pd.to_datetime(observation_period_end) if monetary_value_col: select_columns.append(monetary_value_col) @@ -182,7 +184,7 @@ def _find_first_transactions( transactions.set_index(datetime_col).to_period(time_unit).to_timestamp() ) - mask = pd.DatetimeIndex(transactions.index) <= observation_period_end + mask = pd.to_datetime(transactions.index) <= pd.to_datetime(observation_period_end) transactions = transactions.loc[mask].reset_index()