From 84c464e3a3c64c0e98d232acc9c2e19fd58ce66f Mon Sep 17 00:00:00 2001
From: SangGyu An <bbeat2782@gmail.com>
Date: Wed, 28 Jun 2023 16:35:31 -0700
Subject: [PATCH] Making ResultSet distinguishable (#655)

---
 CHANGELOG.md                |  2 ++
 src/sql/run.py              | 22 ++++++++++++++++++++++
 src/tests/test_resultset.py | 28 +++++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fb69fe39a..b70755878 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 * [Feature] Modified `TableDescription` to add styling, generate messages and format the calculated outputs (#459)
 * [Feature] Support flexible spacing `myvar=<<` operator ([#525](https://github.com/ploomber/jupysql/issues/525))
+* [Feature] Added a line under `ResultSet` to distinguish it from data frame and error message when invalid operations are performed (#468)
+
 * [Doc] Modified integrations content to ensure they're all consistent (#523)
 * [Doc] Document --persist-replace in API section (#539)
 * [Fix] Fixed CI issue by updating `invalid_connection_string_duckdb` in `test_magic.py` (#631)
diff --git a/src/sql/run.py b/src/sql/run.py
index fb8bcaff3..e6e76cd34 100644
--- a/src/sql/run.py
+++ b/src/sql/run.py
@@ -126,6 +126,16 @@ def _repr_html_(self):
         if self.pretty:
             self.pretty.add_rows(self)
             result = self.pretty.get_html_string()
+            HTML = (
+                "%s\n<span style='font-style:italic;font-size:11px'>"
+                "<code>ResultSet</code> : to convert to pandas, call <a href="
+                "'https://jupysql.ploomber.io/en/latest/integrations/pandas.html'>"
+                "<code>.DataFrame()</code></a> or to polars, call <a href="
+                "'https://jupysql.ploomber.io/en/latest/integrations/polars.html'>"
+                "<code>.PolarsDataFrame()</code></a></span><br>"
+            )
+            result = HTML % (result)
+
             # to create clickable links
             result = html.unescape(result)
             result = _cell_with_spaces_pattern.sub(_nonbreaking_spaces, result)
@@ -179,6 +189,18 @@ def __getitem__(self, key):
                 raise KeyError('%d results for "%s"' % (len(result), key))
             return result[0]
 
+    def __getattribute__(self, attr):
+        "Raises AttributeError when invalid operation is performed."
+        try:
+            return object.__getattribute__(self, attr)
+        except AttributeError:
+            err_msg = (
+                f"'{attr}' is not a valid operation, you can convert this "
+                "into a pandas data frame by calling '.DataFrame()' or a "
+                "polars data frame by calling '.PolarsDataFrame()'"
+            )
+            raise AttributeError(err_msg) from None
+
     def dict(self):
         """Returns a single dict built from the result set
 
diff --git a/src/tests/test_resultset.py b/src/tests/test_resultset.py
index 52c7f338e..4bf6cb337 100644
--- a/src/tests/test_resultset.py
+++ b/src/tests/test_resultset.py
@@ -77,7 +77,33 @@ def test_resultset_repr_html(result_set):
         "<th>x</th>\n        </tr>\n    </thead>\n    <tbody>\n        "
         "<tr>\n            <td>0</td>\n        </tr>\n        <tr>\n            "
         "<td>1</td>\n        </tr>\n        <tr>\n            <td>2</td>\n        "
-        "</tr>\n    </tbody>\n</table>"
+        "</tr>\n    </tbody>\n</table>\n"
+        "<span style='font-style:italic;font-size:11px'>"
+        "<code>ResultSet</code> : to convert to pandas, call <a href="
+        "'https://jupysql.ploomber.io/en/latest/integrations/pandas.html'>"
+        "<code>.DataFrame()</code></a> or to polars, call <a href="
+        "'https://jupysql.ploomber.io/en/latest/integrations/polars.html'>"
+        "<code>.PolarsDataFrame()</code></a></span><br>"
+    )
+
+
+@pytest.mark.parametrize(
+    "fname, parameters",
+    [
+        ("head", -1),
+        ("tail", None),
+        ("value_counts", None),
+        ("not_df_function", None),
+    ],
+)
+def test_invalid_operation_error(result_set, fname, parameters):
+    with pytest.raises(AttributeError) as excinfo:
+        getattr(result_set, fname)(parameters)
+
+    assert str(excinfo.value) == (
+        f"'{fname}' is not a valid operation, you can convert this "
+        "into a pandas data frame by calling '.DataFrame()' or a "
+        "polars data frame by calling '.PolarsDataFrame()'"
     )