Fix/highlight correct item (#11016)

* get children of unexpanded books

* highlight item inside collapsed parent in treeView

* fix path issue on windows

* refactor common code into separate func

* refactor

extensions/notebook/src/book/bookTreeView.ts

@@ -180,7 +180,10 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 		}
 		this._bookViewer = this._apiWrapper.createTreeView(this.viewId, { showCollapseAll: true, treeDataProvider: this });
 		this._bookViewer.onDidChangeVisibility(e => {
-			if (e.visible) {
+			let openDocument = azdata.nb.activeNotebookEditor;
+			let notebookPath = openDocument?.document.uri;
+			// call reveal only once on the correct view
+			if (e.visible && ((!this._openAsUntitled && notebookPath?.scheme !== 'untitled') || (this._openAsUntitled && notebookPath?.scheme === 'untitled'))) {
 				this.revealActiveDocumentInViewlet();
 			}
 		});
@@ -230,24 +233,49 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 
 	async revealActiveDocumentInViewlet(uri?: vscode.Uri, shouldReveal: boolean = true): Promise<void> {
 		let bookItem: BookTreeItem;
+		let notebookPath: string;
 		// If no uri is passed in, try to use the current active notebook editor
 		if (!uri) {
 			let openDocument = azdata.nb.activeNotebookEditor;
 			if (openDocument) {
-				bookItem = this.currentBook?.getNotebook(openDocument.document.uri.fsPath);
+				notebookPath = openDocument.document.uri.fsPath;
 			}
 		} else if (uri.fsPath) {
-			bookItem = this.currentBook?.getNotebook(uri.fsPath);
+			notebookPath = uri.fsPath;
 		}
+		bookItem = await this.findAndExpandParentNode(notebookPath);
+
 		if (bookItem) {
 			// Select + focus item in viewlet if books viewlet is already open, or if we pass in variable
 			if (shouldReveal || this._bookViewer.visible) {
 				// Note: 3 is the maximum number of levels that the vscode APIs let you expand to
-				await this._bookViewer.reveal(bookItem, { select: true, focus: true, expand: 3 });
+				await this._bookViewer.reveal(bookItem, { select: true, focus: true, expand: true });
 			}
 		}
 	}
 
+	async findAndExpandParentNode(notebookPath: string): Promise<BookTreeItem> {
+		let bookItem: BookTreeItem = this.currentBook?.getNotebook(notebookPath);
+		// if the node is not expanded getNotebook returns undefined, try to expand the parent node or getChildren of
+		// the root node.
+		if (!bookItem) {
+			// get the parent node and expand it if it's not already
+			let allNodes = this.currentBook?.getAllNotebooks();
+			let book = allNodes ? Array.from(allNodes?.keys())?.filter(x => x.indexOf(notebookPath.substring(0, notebookPath.lastIndexOf(path.sep))) > -1) : undefined;
+			let bookNode = book?.length > 0 ? this.currentBook?.getNotebook(book.find(x => x.substring(0, x.lastIndexOf(path.sep)) === notebookPath.substring(0, notebookPath.lastIndexOf(path.sep)))) : undefined;
+			if (bookNode) {
+				if (this._bookViewer.visible) {
+					await this._bookViewer.reveal(bookNode, { select: true, focus: false, expand: 3 });
+				} else {
+					await this.getChildren(bookNode);
+				}
+
+				bookItem = this.currentBook?.getNotebook(notebookPath);
+			}
+		}
+		return bookItem;
+	}
+
 	openMarkdown(resource: string): void {
 		this.runThrottledAction(resource, () => {
 			try {
@@ -459,21 +487,13 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 
 	getParent(element?: BookTreeItem): vscode.ProviderResult<BookTreeItem> {
 		if (element) {
-			let parentPath;
+			let parentPath: string;
-			if (element.root.endsWith('.md')) {
+			parentPath = path.join(element.root, Content, element.uri.substring(0, element.uri.lastIndexOf(path.posix.sep)));
-				parentPath = path.join(this.currentBook.bookPath, Content, 'readme.md');
+			if (parentPath === element.root) {
-				if (parentPath === element.root) {
-					return undefined;
-				}
-			}
-			else if (element.root.endsWith('.ipynb')) {
-				let baseName: string = path.basename(element.root);
-				parentPath = element.root.replace(baseName, 'readme.md');
-			}
-			else {
 				return undefined;
 			}
-			return this.currentBook.getAllNotebooks().get(parentPath);
+			let parentPaths = Array.from(this.currentBook.getAllNotebooks()?.keys()).filter(x => x.indexOf(parentPath) > -1);
+			return parentPaths.length > 0 ? this.currentBook.getAllNotebooks().get(parentPaths[0]) : undefined;
 		} else {
 			return undefined;
 		}

extensions/notebook/src/extension.ts

@@ -15,7 +15,7 @@ import { IExtensionApi, IPackageManageProvider } from './types';
 import { CellType } from './contracts/content';
 import { NotebookUriHandler } from './protocol/notebookUriHandler';
 import { BookTreeViewProvider } from './book/bookTreeView';
-import { NavigationProviders } from './common/constants';
+import { NavigationProviders, BuiltInCommands, unsavedBooksContextKey } from './common/constants';
 
 const localize = nls.loadMessageBundle();
 
@@ -128,7 +128,14 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
 		} else {
 			bookTreeViewProvider.revealActiveDocumentInViewlet(e.document.uri, false);
 		}
+	});
 
+	azdata.nb.onDidOpenNotebookDocument(async e => {
+		if (e.uri.scheme === 'untitled') {
+			await vscode.commands.executeCommand(BuiltInCommands.SetContext, unsavedBooksContextKey, true);
+		} else {
+			await vscode.commands.executeCommand(BuiltInCommands.SetContext, unsavedBooksContextKey, false);
+		}
 	});
 
 	return {