fix: refactor the code

Author: mayurkmmt

pyretailscience/plots/index.py

@@ -284,16 +284,15 @@ def get_indexes(
 
     if index_subgroup_col is None:
         overall_total = overall_agg.value.sum().execute()
-        overall_props = overall_agg.mutate(proportion=overall_agg.value / overall_total)
+        overall_props = overall_agg.mutate(proportion_overall=overall_agg.value / overall_total)
     else:
         overall_total = overall_agg.group_by(index_subgroup_col).aggregate(total=lambda t: t.value.sum())
         overall_props = (
             overall_agg.join(overall_total, index_subgroup_col)
-            .mutate(proportion=lambda t: t.value / t.total)
+            .mutate(proportion_overall=lambda t: t.value / t.total)
             .drop("total")
         )
 
-    overall_props = overall_props.mutate(proportion_overall=overall_props.proportion).drop("proportion")
     table = table.filter(table[index_col] == value_to_index)
     subset_agg = table.group_by(group_cols).aggregate(value=agg_fn(table[value_col]))
 

tests/plots/test_index.py

@@ -1,5 +1,6 @@
 """Tests for the index plot module."""
 
+import ibis
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -8,7 +9,7 @@
 from pyretailscience.plots.index import get_indexes, plot
 
 OFFSET_VALUE = 100
-OFFSET_THRESHOLD = -5
+OFFSET_THRESHOLD = 5
 
 
 def test_get_indexes_basic():
@@ -109,14 +110,14 @@ def test_get_indexes_with_offset():
         index_col="category",
         value_col="value",
         group_col="category",
-        offset=5,
+        offset=OFFSET_THRESHOLD,
     )
 
     assert isinstance(result, pd.DataFrame)
     assert "category" in result.columns
     assert "index" in result.columns
     assert not result.empty
-    assert all(result["index"] >= OFFSET_THRESHOLD)
+    assert all(result["index"] >= -OFFSET_THRESHOLD)
 
 
 def test_get_indexes_single_column():
@@ -168,6 +169,23 @@ def test_get_indexes_two_columns():
     pd.testing.assert_frame_equal(output, expected_output)
 
 
+def test_get_indexes_with_ibis_table_input():
+    """Test that the get_indexes function works with an ibis Table."""
+    df = pd.DataFrame(
+        {
+            "category": ["A", "B", "C"],
+            "value": [10, 20, 30],
+        },
+    )
+    table = ibis.memtable(df)
+
+    result = get_indexes(table, value_to_index="A", index_col="category", value_col="value", group_col="category")
+    assert isinstance(result, pd.DataFrame)
+    assert "category" in result.columns
+    assert "index" in result.columns
+    assert not result.empty
+
+
 class TestIndexPlot:
     """Tests for the index_plot function."""