Merge pull request #103 from e10v/dev

Delete `to_ibis` parameter from `make_users_data` and `make_sessions_data`
e10v · Dec 14, 2024 · 69f738c · 69f738c
2 parents 9897585 + d67d0d7
commit 69f738c
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 219 deletions.
diff --git a/src/tea_tasting/datasets.py b/src/tea_tasting/datasets.py
@@ -3,23 +3,20 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, overload
+from typing import TYPE_CHECKING
 
-import ibis
 import numpy as np
 import pandas as pd
 
 import tea_tasting.utils
 
 
 if TYPE_CHECKING:
-    from typing import Any, Literal
+    from typing import Any
 
-    import ibis.expr.types  # noqa: TC004
     import numpy.typing as npt
 
 
-@overload
 def make_users_data(
     *,
     covariates: bool = False,
@@ -32,41 +29,7 @@ def make_users_data(
     avg_sessions: float | int = 2,
     avg_orders_per_session: float = 0.25,
     avg_revenue_per_order: float | int = 10,
-    to_ibis: Literal[False] = False,
 ) -> pd.DataFrame:
-    ...
-
-@overload
-def make_users_data(
-    *,
-    covariates: bool = False,
-    seed: int | np.random.Generator | np.random.SeedSequence | None = None,
-    n_users: int = 4000,
-    ratio: float | int = 1,
-    sessions_uplift: float | int = 0.0,
-    orders_uplift: float = 0.1,
-    revenue_uplift: float = 0.1,
-    avg_sessions: float | int = 2,
-    avg_orders_per_session: float = 0.25,
-    avg_revenue_per_order: float | int = 10,
-    to_ibis: Literal[True] = True,
-) -> ibis.expr.types.Table:
-    ...
-
-def make_users_data(
-    *,
-    covariates: bool = False,
-    seed: int | np.random.Generator | np.random.SeedSequence | None = None,
-    n_users: int = 4000,
-    ratio: float | int = 1,
-    sessions_uplift: float | int = 0.0,
-    orders_uplift: float = 0.1,
-    revenue_uplift: float = 0.1,
-    avg_sessions: float | int = 2,
-    avg_orders_per_session: float = 0.25,
-    avg_revenue_per_order: float | int = 10,
-    to_ibis: bool = False,
-) -> ibis.expr.types.Table | pd.DataFrame:
     """Generate simulated data for A/B testing scenarios.
 
     Data mimics what you might encounter in an A/B test for an online store,
@@ -100,8 +63,6 @@ def make_users_data(
         avg_orders_per_session: Average number of orders per session.
             Should be less than `1`.
         avg_revenue_per_order: Average revenue per order.
-        to_ibis: If set to `True`, returns an Ibis Table; otherwise,
-            returns a Pandas DataFrame.
 
     Returns:
         Simulated data for A/B testing scenarios.
@@ -161,12 +122,10 @@ def make_users_data(
         avg_sessions=avg_sessions,
         avg_orders_per_session=avg_orders_per_session,
         avg_revenue_per_order=avg_revenue_per_order,
-        to_ibis=to_ibis,
         explode_sessions=False,
     )
 
 
-@overload
 def make_sessions_data(
     *,
     covariates: bool = False,
@@ -179,41 +138,7 @@ def make_sessions_data(
     avg_sessions: float | int = 2,
     avg_orders_per_session: float = 0.25,
     avg_revenue_per_order: float | int = 10,
-    to_ibis: Literal[False] = False,
 ) -> pd.DataFrame:
-    ...
-
-@overload
-def make_sessions_data(
-    *,
-    covariates: bool = False,
-    seed: int | np.random.Generator | np.random.SeedSequence | None = None,
-    n_users: int = 4000,
-    ratio: float | int = 1,
-    sessions_uplift: float | int = 0.0,
-    orders_uplift: float = 0.1,
-    revenue_uplift: float = 0.1,
-    avg_sessions: float | int = 2,
-    avg_orders_per_session: float = 0.25,
-    avg_revenue_per_order: float | int = 10,
-    to_ibis: Literal[True] = True,
-) -> ibis.expr.types.Table:
-    ...
-
-def make_sessions_data(
-    *,
-    covariates: bool = False,
-    seed: int | np.random.Generator | np.random.SeedSequence | None = None,
-    n_users: int = 4000,
-    ratio: float | int = 1,
-    sessions_uplift: float | int = 0.0,
-    orders_uplift: float = 0.1,
-    revenue_uplift: float = 0.1,
-    avg_sessions: float | int = 2,
-    avg_orders_per_session: float = 0.25,
-    avg_revenue_per_order: float | int = 10,
-    to_ibis: bool = False,
-) -> ibis.expr.types.Table | pd.DataFrame:
     """Generate simulated user data for A/B testing scenarios.
 
     Data mimics what you might encounter in an A/B test for an online store,
@@ -247,8 +172,6 @@ def make_sessions_data(
         avg_orders_per_session: Average number of orders per session.
             Should be less than `1`.
         avg_revenue_per_order: Average revenue per order.
-        to_ibis: If set to `True`, returns an Ibis Table; otherwise,
-            returns a Pandas DataFrame.
 
     Returns:
         Simulated data for A/B testing scenarios.
@@ -308,7 +231,6 @@ def make_sessions_data(
         avg_sessions=avg_sessions,
         avg_orders_per_session=avg_orders_per_session,
         avg_revenue_per_order=avg_revenue_per_order,
-        to_ibis=to_ibis,
         explode_sessions=True,
     )
 
@@ -325,9 +247,8 @@ def _make_data(
     avg_sessions: float | int = 2,
     avg_orders_per_session: float = 0.25,
     avg_revenue_per_order: float | int = 10,
-    to_ibis: bool = False,
     explode_sessions: bool = False,
-) -> ibis.expr.types.Table | pd.DataFrame:
+) -> pd.DataFrame:
     _check_params(
         n_users=n_users,
         ratio=ratio,
@@ -418,9 +339,6 @@ def _make_data(
             revenue_covariate=revenue_covariate,
         )
 
-    if to_ibis:
-        return ibis.memtable(data)
-
     return data