From 11b74fe7715e0543210b3762aa5b0c9051834a13 Mon Sep 17 00:00:00 2001
From: Pete Blois <blois@google.com>
Date: Tue, 19 Sep 2023 09:20:20 -0700
Subject: [PATCH] No public description

PiperOrigin-RevId: 566653529
---
 google/colab/_reprs.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/google/colab/_reprs.py b/google/colab/_reprs.py
index fe1d8951..1f6c6cf6 100644
--- a/google/colab/_reprs.py
+++ b/google/colab/_reprs.py
@@ -1,5 +1,6 @@
 """Rich representations of built-in types."""
 
+import re
 import warnings
 # pytype: disable=import-error
 import IPython
@@ -145,10 +146,24 @@ def _dataframe_intrinsic_repr(dataframe):
   }
   if ip := IPython.get_ipython():
     namespace = ip.user_ns
+    found = False
     for varname, var in namespace.items():
       if dataframe is var and not varname.startswith('_'):
         result['variable_name'] = varname
+        found = True
         break
+    if not found:
+      last_line = (
+          ip.parent_header.get('content', {}).get('code', '').split('\n')[-1]
+      )
+      if match := re.match(r'^([a-z_A-Z0-9]+).head(\s?[0-9]*\s?)', last_line):
+        import pandas as pd
+
+        varname = match.groups()[0]
+        possible_df = ip.user_ns.get(varname)
+        if isinstance(possible_df, pd.DataFrame):
+          result['variable_name'] = varname
+          dataframe = possible_df
 
   if summary := _summarize_dataframe(dataframe):
     result['summary'] = summary