From f4593d5d79d68f1b070157127b894bf667a37508 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 10 Dec 2024 10:05:53 -0800
Subject: [PATCH] Highlight XML strings using pygments in mjSpec notebook.

PiperOrigin-RevId: 704751058
Change-Id: Iaf7fae7f13ef12d59bac762170a6dfca01494164
---
 python/mjspec.ipynb | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb
index f6fa9fe9..c1e1fbf9 100644
--- a/python/mjspec.ipynb
+++ b/python/mjspec.ipynb
@@ -109,7 +109,7 @@
     "# More legible printing from numpy.\n",
     "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
     "\n",
-    "from IPython.display import clear_output\n",
+    "from IPython.display import clear_output, HTML, display\n",
     "clear_output()\n",
     "\n",
     "# Get MuJoCo's humanoid model and a Franka arm from the MuJoCo Menagerie.\n",
@@ -121,6 +121,12 @@
     "!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
     "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
     "\n",
+    "def print_xml(xml_string):\n",
+    "  formatter = pygments.formatters.HtmlFormatter(style='lovelace')\n",
+    "  lexer = pygments.lexers.XmlLexer()\n",
+    "  highlighted = pygments.highlight(xml_string, lexer, formatter)\n",
+    "  display(HTML(f\"<style>{formatter.get_style_defs()}</style>{highlighted}\"))\n",
+    "\n",
     "def render(model, data=None, height=250):\n",
     "  if data is None:\n",
     "    data = mj.MjData(model)\n",
@@ -172,6 +178,7 @@
     "render(model)\n",
     "\n",
     "# Change the mjSpec, re-compile and re-render\n",
+    "spec.modelname = \"edited model\"\n",
     "geoms = spec.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n",
     "geoms[0].name = 'blue_box'\n",
     "geoms[0].rgba = [0, 0, 1, 1]\n",
@@ -193,7 +200,7 @@
     "id": "Tw_yUwqxKwCI"
    },
    "source": [
-    "`mjSpec` can save XML to string, saving all modifications."
+    "`mjSpec` can save XML to string, with all modifications:"
    ]
   },
   {
@@ -204,9 +211,7 @@
    },
    "outputs": [],
    "source": [
-    "#@title Print an XML from an `mjSpec` {vertical-output: true}\n",
-    "\n",
-    "print(spec.to_xml())"
+    "print_xml(spec.to_xml())"
    ]
   },
   {