From 6b46de21a7e1f38993512fad910fbce4e7b9577d Mon Sep 17 00:00:00 2001
From: Kasra Ferdowsifard <kferdowsifard@gmail.com>
Date: Wed, 22 Mar 2023 17:21:48 -0700
Subject: [PATCH] fixed bug where plt.show() showed multiple plots

---
 .../editor/contrib/rtv/browser/RTVDisplay.ts  | 55 ++++++++++++++++---
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/vs/editor/contrib/rtv/browser/RTVDisplay.ts b/src/vs/editor/contrib/rtv/browser/RTVDisplay.ts
index 4f4ff6e08184a..11c6ffed228c1 100644
--- a/src/vs/editor/contrib/rtv/browser/RTVDisplay.ts
+++ b/src/vs/editor/contrib/rtv/browser/RTVDisplay.ts
@@ -236,6 +236,8 @@ class RTVOutputDisplayBox {
 	private _box: HTMLDivElement;
 	private _html: string = '<b>Output:</b><br><br><b>Errors:</b><br>';
 	private _isOnDiv: boolean = false;
+	private _outOfDate: boolean = false;
+
 	constructor(
 		private readonly _editor: ICodeEditor
 	) {
@@ -332,7 +334,12 @@ class RTVOutputDisplayBox {
 		return this._isOnDiv;
 	}
 
+	public outOfDate() {
+		this._outOfDate = true;
+	}
+
 	public update(outputMsg: string, errorMsg: string, parsedResults: any) {
+		this._outOfDate = false;
 
 		function escapeHTML(unsafe: string): string {
 			return unsafe
@@ -412,16 +419,46 @@ class RTVOutputDisplayBox {
 							continue;
 						}
 
-						for (const map of env) {
-							if (!('Plot' in map)) {
-								console.error('No entry named "Plot" found in map for line ', lineNumber, '. Skipping.\nmap:', map, '\nenv:\n', env, '\nenvs:\n', envs);
-								continue;
+						// We need to check how many plots we need to show.
+						// Basically, if plt.show() was in a loop, then its env will have multiple maps.
+						// If it was called outside a loop however, it should only have one.
+						// We care about this because of the cornercase where the last actual Plot was
+						// in a loop, but plt.show() was called once outside the loop.
+						// TODO This is a really bad hack, and will break if, e.g., the last Plot and
+						// plt.show() are called in _separate_ loops.
+
+						const showEnv = envs[lineNumber];
+						if (showEnv.length === 1) {
+							// Only print the _last_ plot.
+							for (const map of env.reverse()) {
+								if ('Plot' in map) {
+									// Clean up the plot!
+									let plot: string = map['Plot'];
+									plot = plot.substring(8, plot.length - 3);
+									plots.push(plot);
+									break;
+								}
 							}
+						} else {
+							if (showEnv.length !== env.length) {
+								console.error('plt.show() was called ',
+									showEnv.length, ' times on line ',
+									lineNumber, ' but found ',
+									env.length, ' plots before it.');
+							}
+
+							// Print all of them!
+							for (const map of env) {
+								if (!('Plot' in map)) {
+									console.error('No entry named "Plot" found in map for line ', lineNumber, '. Skipping.\nmap:', map, '\nenv:\n', env, '\nenvs:\n', envs);
+									continue;
+								}
 
-							// Clean up the plot!
-							let plot: string = map['Plot'];
-							plot = plot.substring(8, plot.length - 3);
-							plots.push(plot);
+								// Clean up the plot!
+								let plot: string = map['Plot'];
+								plot = plot.substring(8, plot.length - 3);
+								plots.push(plot);
+							}
 						}
 					}
 				}
@@ -429,7 +466,7 @@ class RTVOutputDisplayBox {
 
 			if (plots.length > 0) {
 				// Show the plot as well.
-				const plotsHtml = plots.join('\n<br>\n');
+				const plotsHtml = plots.join('\n<p>\n');
 				this.setContent(`<b>Output:</b><pre>${escapeHTML(outputMsg)}</pre><b>Errors:</b><pre>${errorMsgStyled}</pre><b>Plots:</b><br>${plotsHtml}`);
 			} else {
 				// Only Output and Error.