Implement the ultimate quick pick experience

utils/src/tree/AzExtParentTreeItem.ts

@@ -246,7 +246,7 @@ export abstract class AzExtParentTreeItem extends AzExtTreeItem implements types
                 // Just in case implementers of `loadMoreChildrenImpl` re-use the same child node, we want to clear those caches as well
                 for (const child of this._cachedChildren) {
                     if (isAzExtParentTreeItem(child)) {
-                        (<AzExtParentTreeItem>child).clearCache();
+                        child.clearCache();
                     }
                 }
                 this._cachedChildren = [];

utils/src/tree/AzExtTreeDataProvider.ts

@@ -147,7 +147,7 @@ export class AzExtTreeDataProvider implements IAzExtTreeDataProviderInternal, ty
         }
 
         if (isAzExtParentTreeItem(treeItem)) {
-            (<AzExtParentTreeItem>treeItem).clearCache();
+            treeItem.clearCache();
         }
 
         this.refreshUIOnly(treeItem);
@@ -179,7 +179,7 @@ export class AzExtTreeDataProvider implements IAzExtTreeDataProviderInternal, ty
 
         while (!treeItem.matchesContextValue(expectedContextValues)) {
             if (isAzExtParentTreeItem(treeItem)) {
-                const pickedItems: AzExtTreeItem | AzExtTreeItem[] = await (<AzExtParentTreeItem>treeItem).pickChildTreeItem(expectedContextValues, context);
+                const pickedItems: AzExtTreeItem | AzExtTreeItem[] = await treeItem.pickChildTreeItem(expectedContextValues, context);
                 if (Array.isArray(pickedItems)) {
                     // canPickMany is only supported at the last stage of the picker, so automatically return if this is an array
                     return <T[]><unknown>pickedItems;

utils/src/tree/InternalInterfaces.ts

@@ -5,13 +5,14 @@
 
 import { EventEmitter } from 'vscode';
 import * as types from '../../index';
+import { AzExtParentTreeItem } from './AzExtParentTreeItem';
 import { AzExtTreeItem } from './AzExtTreeItem';
 import { CollapsibleStateTracker } from './CollapsibleStateTracker';
 
 // Interfaces for methods on the tree that aren't exposed outside of this package
 // We can't reference the classes directly because it would result in circular dependencies
 
-export interface IAzExtParentTreeItemInternal extends types.AzExtParentTreeItem, AzExtTreeItem {
+export interface IAzExtParentTreeItemInternal extends AzExtParentTreeItem {
     _isAzExtParentTreeItem: boolean;
     parent: IAzExtParentTreeItemInternal | undefined;
     treeDataProvider: IAzExtTreeDataProviderInternal;
@@ -28,6 +29,6 @@ export interface IAzExtTreeDataProviderInternal extends types.AzExtTreeDataProvi
 /**
  * Using instanceof AzExtParentTreeItem causes issues whenever packages are linked for dev testing. Instead, check _isAzExtParentTreeItem
  */
-export function isAzExtParentTreeItem(item: {}): boolean {
-    return !!(<IAzExtParentTreeItemInternal>item)._isAzExtParentTreeItem;
+export function isAzExtParentTreeItem(maybeParent: unknown): maybeParent is IAzExtParentTreeItemInternal {
+    return !!(maybeParent as IAzExtParentTreeItemInternal)._isAzExtParentTreeItem;
 }

utils/src/treev2/quickPickWizard/ContextValueQuickPickStep.ts

@@ -0,0 +1,89 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../../index';
+import { GenericQuickPickOptions, GenericQuickPickStep } from './GenericQuickPickStep';
+import { isAzExtParentTreeItem } from '../../tree/InternalInterfaces';
+import { QuickPickWizardContext } from './QuickPickWizardContext';
+
+/**
+ * Describes filtering based on context value. Items that pass the filter will
+ * match at least one of the `include` filters, but none of the `exclude` filters.
+ */
+export interface ContextValueFilter {
+    /**
+     * This filter will include items that match *any* of the values in the array.
+     * When a string is used, exact value comparison is done.
+     */
+    include: string | RegExp | (string | RegExp)[];
+
+    /**
+     * This filter will exclude items that match *any* of the values in the array.
+     * When a string is used, exact value comparison is done.
+     */
+    exclude?: string | RegExp | (string | RegExp)[];
+}
+
+interface ContextValueFilterableTreeNodeV2 {
+    readonly quickPickOptions: {
+        readonly contextValues: Array<string>;
+        readonly isLeaf: boolean;
+    }
+}
+
+export type ContextValueFilterableTreeNode = ContextValueFilterableTreeNodeV2 | types.AzExtTreeItem;
+
+export interface ContextValueFilterQuickPickOptions extends GenericQuickPickOptions {
+    contextValueFilter: ContextValueFilter;
+}
+
+export class ContextValueQuickPickStep<TNode extends ContextValueFilterableTreeNode, TContext extends QuickPickWizardContext<TNode>, TOptions extends ContextValueFilterQuickPickOptions> extends GenericQuickPickStep<TNode, TContext, TOptions> {
+    protected override isDirectPick(node: TNode): boolean {
+        const includeOption = this.pickOptions.contextValueFilter.include;
+        const excludeOption = this.pickOptions.contextValueFilter.exclude;
+
+        const includeArray: (string | RegExp)[] = Array.isArray(includeOption) ? includeOption : [includeOption];
+        const excludeArray: (string | RegExp)[] = excludeOption ?
+            (Array.isArray(excludeOption) ? excludeOption : [excludeOption]) :
+            [];
+
+        const nodeContextValues: string[] = isV2TreeNode(node) ?
+            node.quickPickOptions.contextValues :
+            [node.contextValue];
+
+        return includeArray.some(i => this.matchesSingleFilter(i, nodeContextValues)) &&
+            !excludeArray.some(e => this.matchesSingleFilter(e, nodeContextValues));
+    }
+
+    protected override isIndirectPick(node: TNode): boolean {
+        if (isV2TreeNode(node)) {
+            return node.quickPickOptions.isLeaf === false;
+        } else if (isAzExtParentTreeItem(node)) {
+            return true;
+        }
+
+        return false;
+    }
+
+    private matchesSingleFilter(matcher: string | RegExp, nodeContextValues: string[]): boolean {
+        return nodeContextValues.some(c => {
+            if (matcher instanceof RegExp) {
+                return matcher.test(c);
+            }
+
+            // Context value matcher is a string, do full equality (same as old behavior)
+            return c === matcher;
+        })
+    }
+}
+
+function isV2TreeNode(maybeNode: unknown): maybeNode is ContextValueFilterableTreeNodeV2 {
+    if (typeof maybeNode === 'object') {
+        return Array.isArray((maybeNode as ContextValueFilterableTreeNodeV2).quickPickOptions?.contextValues) &&
+            (maybeNode as ContextValueFilterableTreeNodeV2).quickPickOptions?.isLeaf !== null;
+    }
+
+    return false;
+}

utils/src/treev2/quickPickWizard/GenericQuickPickStep.ts

@@ -0,0 +1,112 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../../index';
+import * as vscode from 'vscode';
+import { getLastNode, QuickPickWizardContext } from './QuickPickWizardContext';
+import { AzureWizardPromptStep } from '../../wizard/AzureWizardPromptStep';
+import { NoResourceFoundError } from '../../errors';
+import { parseError } from '../../parseError';
+
+export interface GenericQuickPickOptions {
+    skipIfOne?: boolean;
+}
+
+export interface SkipIfOneQuickPickOptions extends GenericQuickPickOptions {
+    skipIfOne?: true;
+}
+
+export abstract class GenericQuickPickStep<TNode extends unknown, TContext extends QuickPickWizardContext<TNode>, TOptions extends GenericQuickPickOptions> extends AzureWizardPromptStep<TContext> {
+    public readonly supportsDuplicateSteps = true;
+
+    public constructor(
+        protected readonly treeDataProvider: vscode.TreeDataProvider<TNode>,
+        protected readonly pickOptions: TOptions
+    ) {
+        super();
+    }
+
+    public async prompt(wizardContext: TContext): Promise<void> {
+        try {
+            const pick = await this.promptInternal(wizardContext);
+            wizardContext.pickedNodes.push(pick);
+        } catch (err) {
+            const error = parseError(err);
+            if (error.errorType === 'GoBackError') {
+                // Instead of wiping out a property value, which is the default wizard behavior for `GoBackError`, pop the most recent
+                // value off from the provenance of the picks
+                wizardContext.pickedNodes.pop();
+            }
+
+            // And rethrow
+            throw err;
+        }
+    }
+
+    public shouldPrompt(_wizardContext: TContext): boolean {
+        return true;
+    }
+
+    protected async promptInternal(wizardContext: TContext): Promise<TNode> {
+        const picks = await this.getPicks(wizardContext);
+
+        if (picks.length === 1 && this.pickOptions.skipIfOne) {
+            return picks[0].data;
+        } else {
+            const selected = await wizardContext.ui.showQuickPick(picks, { /* TODO: options */ });
+            return selected.data;
+        }
+    }
+
+    protected async getPicks(wizardContext: TContext): Promise<types.IAzureQuickPickItem<TNode>[]> {
+        const lastPickedItem: TNode | undefined = getLastNode(wizardContext);
+
+        // TODO: if `lastPickedItem` is an `AzExtParentTreeItem`, should we clear its cache?
+        const children = (await this.treeDataProvider.getChildren(lastPickedItem)) || [];
+
+        const directChoices = children.filter(c => this.isDirectPick(c));
+        const indirectChoices = children.filter(c => this.isIndirectPick(c));
+
+        let promptChoices: TNode[];
+        if (directChoices.length === 0) {
+            if (indirectChoices.length === 0) {
+                throw new NoResourceFoundError();
+            } else {
+                promptChoices = indirectChoices;
+            }
+        } else {
+            promptChoices = directChoices;
+        }
+
+        const picks: types.IAzureQuickPickItem<TNode>[] = [];
+        for (const choice of promptChoices) {
+            picks.push(await this.getQuickPickItem(choice));
+        }
+
+        return picks;
+    }
+
+    /**
+     * Filters for nodes that match the final target.
+     * @param node The node to apply the filter to
+     */
+    protected abstract isDirectPick(node: TNode): boolean;
+
+    /**
+     * Filters for nodes that could have a descendant matching the final target.
+     * @param node The node to apply the filter to
+     */
+    protected abstract isIndirectPick(node: TNode): boolean;
+
+    private async getQuickPickItem(resource: TNode): Promise<types.IAzureQuickPickItem<TNode>> {
+        const treeItem = await Promise.resolve(this.treeDataProvider.getTreeItem(resource));
+
+        return {
+            label: ((treeItem.label as vscode.TreeItemLabel)?.label || treeItem.label) as string,
+            description: treeItem.description as string,
+            data: resource,
+        };
+    }
+}

utils/src/treev2/quickPickWizard/QuickPickWithCreateStep.ts

@@ -0,0 +1,42 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../../index';
+import { getLastNode, QuickPickWizardContext } from './QuickPickWizardContext';
+import { ContextValueFilterableTreeNode, ContextValueFilterQuickPickOptions, ContextValueQuickPickStep } from './ContextValueQuickPickStep';
+import { localize } from '../../localize';
+
+type CreateCallback = <TNode extends ContextValueFilterableTreeNode>() => TNode | Promise<TNode>;
+interface CreateQuickPickOptions extends ContextValueFilterQuickPickOptions {
+    skipIfOne?: never; // Not allowed in CreateQuickPickStep
+    createLabel?: string;
+    createCallback: CreateCallback;
+}
+
+export class CreateQuickPickStep<TNode extends ContextValueFilterableTreeNode, TContext extends QuickPickWizardContext<TNode>> extends ContextValueQuickPickStep<TNode, TContext, CreateQuickPickOptions> {
+    public override async prompt(wizardContext: TContext): Promise<void> {
+        await super.prompt(wizardContext);
+
+        const lastNode = getLastNode(wizardContext) as (TNode | CreateCallback);
+        if (typeof lastNode === 'function') {
+            // If the last node is a function, pop it off the list and execute it
+            const callback = wizardContext.pickedNodes.pop() as unknown as CreateCallback;
+            wizardContext.pickedNodes.push(await callback());
+        }
+    }
+
+    protected override async getPicks(wizardContext: TContext): Promise<types.IAzureQuickPickItem<TNode>[]> {
+        const picks: types.IAzureQuickPickItem<TNode | CreateCallback>[] = await super.getPicks(wizardContext);
+        picks.push(this.getCreatePick());
+        return picks as types.IAzureQuickPickItem<TNode>[];
+    }
+
+    private getCreatePick(): types.IAzureQuickPickItem<CreateCallback> {
+        return {
+            label: this.pickOptions.createLabel || localize('createQuickPickLabel', '$(add) Create...'),
+            data: this.pickOptions.createCallback,
+        };
+    }
+}

utils/src/treev2/quickPickWizard/QuickPickWizardContext.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../../index';
+
+export interface QuickPickWizardContext<TNode extends unknown> extends types.IActionContext {
+    pickedNodes: TNode[];
+}
+
+export function getLastNode<TNode extends unknown>(context: QuickPickWizardContext<TNode>): TNode | undefined {
+    if (context.pickedNodes.length) {
+        return context.pickedNodes[context.pickedNodes.length - 1];
+    }
+
+    return undefined;
+}

utils/src/treev2/quickPickWizard/RecursiveQuickPickStep.ts

@@ -0,0 +1,33 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../../index';
+import { ContextValueFilterableTreeNode, ContextValueFilterQuickPickOptions, ContextValueQuickPickStep } from './ContextValueQuickPickStep';
+import { getLastNode, QuickPickWizardContext } from './QuickPickWizardContext';
+
+export class RecursiveQuickPickStep<TNode extends ContextValueFilterableTreeNode, TContext extends QuickPickWizardContext<TNode>> extends ContextValueQuickPickStep<TNode, TContext, ContextValueFilterQuickPickOptions> {
+    public async getSubWizard(wizardContext: TContext): Promise<types.IWizardOptions<TContext> | undefined> {
+        const lastPickedItem = getLastNode(wizardContext);
+
+        if (!lastPickedItem) {
+            // Something went wrong, no node was chosen
+            throw new Error('No node was set after prompt step.');
+        }
+
+        if (super.isDirectPick(lastPickedItem)) {
+            // The last picked node matches the expected filter
+            // No need to continue prompting
+            return undefined;
+        } else {
+            // Need to keep going because the last picked node is not a match
+            return {
+                hideStepCount: true,
+                promptSteps: [
+                    new RecursiveQuickPickStep(this.treeDataProvider, this.pickOptions)
+                ],
+            };
+        }
+    }
+}