From c2b604c211d0a22e18029fc875b9186a5e05b944 Mon Sep 17 00:00:00 2001
From: Pete Blois <blois@google.com>
Date: Thu, 28 Dec 2023 16:32:20 -0800
Subject: [PATCH] No public description

PiperOrigin-RevId: 594333807
---
 google/colab/_inspector.py | 14 +++----
 google/colab/_reprs.py     | 84 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 7 deletions(-)

diff --git a/google/colab/_inspector.py b/google/colab/_inspector.py
index d10dde44..d5e5340f 100644
--- a/google/colab/_inspector.py
+++ b/google/colab/_inspector.py
@@ -62,7 +62,7 @@
 _UNAVAILABLE_MODULE_NAME = '<unknown>'
 
 
-def _getdoc(obj):
+def getdoc(obj):
   """Custom wrapper for inspect.getdoc.
 
   IPython.core.oinspect.getdoc wraps Python's inspect.getdoc to catch exceptions
@@ -579,7 +579,7 @@ def info(self, obj, oname='', formatter=None, info=None, detail_level=0):
         out['source'] = source
     if 'source' not in out:
       formatter = formatter or (lambda x: x)
-      docstring = formatter(_getdoc(obj) or '<no docstring>')
+      docstring = formatter(getdoc(obj) or '<no docstring>')
       if docstring:
         out['docstring'] = docstring
 
@@ -603,10 +603,10 @@ def info(self, obj, oname='', formatter=None, info=None, detail_level=0):
       # For classes with an __init__, we set init_definition and init_docstring.
       init = getattr(obj, '__init__', None)
       if init:
-        init_docstring = _getdoc(init)
+        init_docstring = getdoc(init)
         if init_docstring and init_docstring != _BASE_INIT_DOC:
           out['init_docstring'] = init_docstring
-        init_def = _get_source_definition(init)
+        init_def = get_source_definition(init)
         if not init_def:
           init_def = self._getdef(init, oname)
         if init_def:
@@ -619,14 +619,14 @@ def info(self, obj, oname='', formatter=None, info=None, detail_level=0):
         if argspec:
           out['argspec'] = argspec
     elif callable(obj):
-      definition = _get_source_definition(obj)
+      definition = get_source_definition(obj)
       if not definition:
         definition = self._getdef(obj, oname)
       if definition:
         out['definition'] = definition
 
       if not oinspect.is_simple_callable(obj):
-        call_docstring = _getdoc(obj.__call__)
+        call_docstring = getdoc(obj.__call__)
         if call_docstring and call_docstring != _BASE_CALL_DOC:
           out['call_docstring'] = call_docstring
 
@@ -646,7 +646,7 @@ def _iscallable(obj):
   )
 
 
-def _get_source_definition(obj):
+def get_source_definition(obj):
   """Get a source representation of the function definition."""
   try:
     obj = _unwrap(obj)
diff --git a/google/colab/_reprs.py b/google/colab/_reprs.py
index e4683cc3..78807892 100644
--- a/google/colab/_reprs.py
+++ b/google/colab/_reprs.py
@@ -1,9 +1,13 @@
 """Rich representations of built-in types."""
 
+import html
 import json
+import types
 import warnings
+from google.colab import _inspector
 # pytype: disable=import-error
 import IPython
+from IPython.core import oinspect
 # pylint: disable=g-import-not-at-top
 with warnings.catch_warnings():
   # Importing via IPython raises a spurious warning, but avoids a version
@@ -11,6 +15,7 @@
   warnings.simplefilter('ignore')
   from IPython.utils import traitlets
 
+
 _original_string_formatters = {}
 
 _original_df_formatters = {}
@@ -205,3 +210,82 @@ def _summarize_dataframe(df, variable_name):
     )
   except Exception:  # pylint: disable=broad-except
     return None
+
+
+def _function_repr(obj):
+  """Renders a function repr."""
+  try:
+    decl = _inspector.get_source_definition(obj)
+    init = getattr(obj, '__init__', None)
+    if not decl and init:
+      decl = _inspector.get_source_definition(init)
+
+    result = (
+        '<div style="max-width:800px; border: 1px solid'
+        ' var(--colab-border-color);">'
+    )
+    if not decl:
+      return
+
+    result += (
+        '<pre style="white-space: initial; background:'
+        ' var(--colab-secondary-surface-color); padding: 8px 12px;'
+        ' border-bottom: 1px solid var(--colab-border-color);">'
+        + decl
+        + '</pre>'
+    )
+
+    filename = oinspect.find_file(obj) or ''
+    docs = _inspector.getdoc(obj) or '<no docstring>'
+    result += (
+        '<pre style="overflow-x: auto; padding: 8px 12px; max-height: 500px;">'
+    )
+    result += (
+        '<a class="filepath" style="display:none"'
+        f' href="#">{html.escape(filename)}</a>'
+    )
+    result += html.escape(docs) + '</pre>'
+    if filename and '<ipython-input' not in filename:
+      line = oinspect.find_source_lines(obj)
+      result += f"""
+      <script>
+      if (google.colab.kernel.accessAllowed && google.colab.files && google.colab.files.view) {{
+        for (const element of document.querySelectorAll('.filepath')) {{
+          element.style.display = 'block'
+          element.onclick = (event) => {{
+            event.preventDefault();
+            event.stopPropagation();
+            google.colab.files.view(element.textContent, {line});
+          }};
+        }}
+      }}
+      </script>
+      """
+    result += '</div>'
+    return result
+  except Exception:  # pylint: disable=broad-except
+    return None
+
+
+def enable_function_repr():
+  """Enables function and class reprs."""
+
+  shell = IPython.get_ipython()
+  if not shell:
+    return
+
+  html_formatter = shell.display_formatter.formatters['text/html']
+  html_formatter.for_type(types.FunctionType, _function_repr)
+  html_formatter.for_type(type, _function_repr)
+
+
+def disable_function_repr():
+  """Disables function and class metadata."""
+
+  shell = IPython.get_ipython()
+  if not shell:
+    return
+
+  html_formatter = shell.display_formatter.formatters['text/html']
+  html_formatter.pop(types.FunctionType)
+  html_formatter.pop(type)