Support pandas' iloc indexer (huggingface#191)

Author: aymeric-roucher

src/smolagents/local_python_executor.py

@@ -685,6 +685,9 @@ def evaluate_subscript(
     if isinstance(value, pd.core.indexing._LocIndexer):
         parent_object = value.obj
         return parent_object.loc[index]
+    if isinstance(value, pd.core.indexing._iLocIndexer):
+        parent_object = value.obj
+        return parent_object.iloc[index]
     if isinstance(value, (pd.DataFrame, pd.Series, np.ndarray)):
         return value[index]
     elif isinstance(value, pd.core.groupby.generic.DataFrameGroupBy):

tests/test_python_interpreter.py

@@ -808,6 +808,7 @@ def test_pandas(self):
         )
         assert np.array_equal(result.values[0], [104, 1])
 
+        # Test groupby
         code = """import pandas as pd
 data = pd.DataFrame.from_dict([
     {"Pclass": 1, "Survived": 1},
@@ -821,6 +822,21 @@ def test_pandas(self):
         )
         assert result.values[1] == 0.5
 
+        # Test loc and iloc
+        code = """import pandas as pd
+data = pd.DataFrame.from_dict([
+    {"Pclass": 1, "Survived": 1},
+    {"Pclass": 2, "Survived": 0},
+    {"Pclass": 2, "Survived": 1}
+])
+survival_rate_biased = data.loc[data['Survived']==1]['Survived'].mean()
+survival_rate_biased = data.loc[data['Survived']==1]['Survived'].mean()
+survival_rate_sorted = data.sort_values(by='Survived', ascending=False).iloc[0]   
+"""
+        result, _ = evaluate_python_code(
+            code, {}, state={}, authorized_imports=["pandas"]
+        )
+
     def test_starred(self):
         code = """
 from math import radians, sin, cos, sqrt, atan2