json indentation improvement (#178)

* Changes to formatter based on if suggestion has data type as object. Then add {} to completion suggestion

* Working json indentation improvement

* handles completions for labels with dynamic schema references

* cursor lands within brackets or quotes

* clean up

* Restore CompletionFormatter.ts with correct implementation from detached HEAD

* context.getEntityType() not used anymore

* fix lint errors, handle top level sections in formatter

* fix lint error

* switch formatter params, break up formatter try block

* resolve merge

* fix lint errors

Author: gemammercado

src/autocomplete/CompletionFormatter.ts

@@ -1,11 +1,23 @@
-import { CompletionItem, CompletionItemKind, CompletionList, InsertTextFormat } from 'vscode-languageserver';
+import {
+    CompletionItem,
+    CompletionItemKind,
+    CompletionList,
+    InsertTextFormat,
+    Range,
+    Position,
+    TextEdit,
+} from 'vscode-languageserver';
 import { Context } from '../context/Context';
 import { ResourceAttributesSet, TopLevelSection, TopLevelSectionsSet } from '../context/ContextType';
+import { Resource } from '../context/semantic/Entity';
+import { EntityType } from '../context/semantic/SemanticTypes';
 import { NodeType } from '../context/syntaxtree/utils/NodeType';
 import { DocumentType } from '../document/Document';
+import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { EditorSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { getIndentationString } from '../utils/IndentationUtils';
+import { RESOURCE_ATTRIBUTE_TYPES } from './CompletionUtils';
 
 export type CompletionItemData = {
     type?: 'object' | 'array' | 'simple';
@@ -17,6 +29,8 @@ export interface ExtendedCompletionItem extends CompletionItem {
 }
 
 export class CompletionFormatter {
+    // In CompletionFormatter class
+
     private static readonly log = LoggerFactory.getLogger(CompletionFormatter);
     private static instance: CompletionFormatter;
 
@@ -38,12 +52,18 @@ export class CompletionFormatter {
         return `{INDENT${numberOfIndents}}`;
     }
 
-    format(completions: CompletionList, context: Context, editorSettings: EditorSettings): CompletionList {
+    format(
+        completions: CompletionList,
+        context: Context,
+        editorSettings: EditorSettings,
+        lineContent?: string,
+        schemaRetriever?: SchemaRetriever,
+    ): CompletionList {
         try {
             const documentType = context.documentType;
-
-            const formattedItems = completions.items.map((item) => this.formatItem(item, documentType, editorSettings));
-
+            const formattedItems = completions.items.map((item) =>
+                this.formatItem(item, documentType, editorSettings, context, lineContent, schemaRetriever),
+            );
             return {
                 ...completions,
                 items: formattedItems,
@@ -58,6 +78,9 @@ export class CompletionFormatter {
         item: CompletionItem,
         documentType: DocumentType,
         editorSettings: EditorSettings,
+        context: Context,
+        lineContent?: string,
+        schemaRetriever?: SchemaRetriever,
     ): CompletionItem {
         const formattedItem = { ...item };
 
@@ -66,19 +89,198 @@ export class CompletionFormatter {
             return formattedItem;
         }
 
+        // Set filterText for ALL items (including snippets) when in JSON with quotes
+        const isInJsonString = documentType === DocumentType.JSON && context.syntaxNode.type === 'string';
+        if (isInJsonString) {
+            formattedItem.filterText = `"${context.text}"`;
+        }
+
         const textToFormat = item.insertText ?? item.label;
 
         if (documentType === DocumentType.JSON) {
-            formattedItem.insertText = this.formatForJson(textToFormat);
+            const result = this.formatForJson(
+                editorSettings,
+                textToFormat,
+                item,
+                context,
+                lineContent,
+                schemaRetriever,
+            );
+            formattedItem.textEdit = TextEdit.replace(result.range, result.text);
+            if (result.isSnippet) {
+                formattedItem.insertTextFormat = InsertTextFormat.Snippet;
+            }
+            delete formattedItem.insertText;
         } else {
             formattedItem.insertText = this.formatForYaml(textToFormat, item, editorSettings);
         }
 
         return formattedItem;
     }
 
-    private formatForJson(label: string): string {
-        return label;
+    private formatForJson(
+        editorSettings: EditorSettings,
+        label: string,
+        item: CompletionItem,
+        context: Context,
+        lineContent?: string,
+        schemaRetriever?: SchemaRetriever,
+    ): { text: string; range: Range; isSnippet: boolean } {
+        const shouldFormat = context.syntaxNode.type === 'string' && !context.isValue() && lineContent;
+
+        const itemData = item.data as CompletionItemData | undefined;
+
+        let formatAsObject = itemData?.type === 'object';
+        let formatAsArray = itemData?.type === 'array';
+        let formatAsString = false;
+
+        if (this.isTopLevelSection(label)) {
+            if (label === String(TopLevelSection.Description)) {
+                formatAsString = true;
+            } else {
+                formatAsObject = true;
+            }
+        }
+        // If type is not in item.data and we have schemaRetriever, look it up from schema
+        if ((!itemData?.type || itemData?.type === 'simple') && schemaRetriever && context.entity) {
+            const propertyType = this.getPropertyTypeFromSchema(schemaRetriever, context, label);
+
+            switch (propertyType) {
+                case 'object': {
+                    formatAsObject = true;
+                    break;
+                }
+                case 'array': {
+                    formatAsArray = true;
+
+                    break;
+                }
+                case 'string': {
+                    formatAsString = true;
+
+                    break;
+                }
+                // No default
+            }
+        }
+
+        const indentation = ' '.repeat(context.startPosition.column);
+        const indentString = getIndentationString(editorSettings, DocumentType.JSON);
+
+        let replacementText = `${indentation}"${label}":`;
+        let isSnippet = false;
+
+        if (shouldFormat) {
+            isSnippet = true;
+            if (formatAsObject) {
+                replacementText = `${indentation}"${label}": {\n${indentation}${indentString}$0\n${indentation}}`;
+            } else if (formatAsArray) {
+                replacementText = `${indentation}"${label}": [\n${indentation}${indentString}$0\n${indentation}]`;
+            } else if (formatAsString) {
+                replacementText = `${indentation}"${label}": "$0"`;
+            }
+        }
+
+        const range = Range.create(
+            Position.create(context.startPosition.row, 0),
+            Position.create(context.endPosition.row, context.endPosition.column + 1),
+        );
+
+        return {
+            text: replacementText,
+            range: range,
+            isSnippet: isSnippet,
+        };
+    }
+
+    /**
+     * Get the type of a property from the CloudFormation schema
+     * @param schemaRetriever - SchemaRetriever instance to get schemas
+     * @param context - Current context with entity and property path information
+     * @param propertyName - Name of the property to look up
+     * @returns The first type found in the schema ('object', 'array', 'string', etc.) or undefined
+     */
+    private getPropertyTypeFromSchema(
+        schemaRetriever: SchemaRetriever,
+        context: Context,
+        propertyName: string,
+    ): string | undefined {
+        let resourceSchema;
+
+        if (ResourceAttributesSet.has(propertyName)) {
+            return RESOURCE_ATTRIBUTE_TYPES[propertyName];
+        }
+
+        const entity = context.entity;
+        if (!entity || context.getEntityType() !== EntityType.Resource) {
+            return undefined;
+        }
+
+        const resourceType = (entity as Resource).Type;
+        if (!resourceType) {
+            return undefined;
+        }
+
+        try {
+            const combinedSchemas = schemaRetriever.getDefault();
+
+            resourceSchema = combinedSchemas.schemas.get(resourceType);
+            if (!resourceSchema) {
+                return undefined;
+            }
+        } catch {
+            return undefined;
+        }
+
+        const propertiesIndex = context.propertyPath.indexOf('Properties');
+        let propertyPath: string[];
+
+        if (propertiesIndex === -1) {
+            propertyPath = [propertyName];
+        } else {
+            const pathAfterProperties = context.propertyPath.slice(propertiesIndex + 1).map(String);
+
+            if (
+                pathAfterProperties.length > 0 &&
+                pathAfterProperties[pathAfterProperties.length - 1] === context.text
+            ) {
+                propertyPath = [...pathAfterProperties.slice(0, -1), propertyName];
+            } else if (pathAfterProperties[pathAfterProperties.length - 1] === propertyName) {
+                propertyPath = pathAfterProperties;
+            } else {
+                propertyPath = [...pathAfterProperties, propertyName];
+            }
+        }
+
+        // Build JSON pointer path using wildcard notation for array indices
+        // CloudFormation schemas use /properties/Tags/*/Key format for array item properties
+        const schemaPath = propertyPath.map((part) => (Number.isNaN(Number(part)) ? part : '*'));
+        const jsonPointerParts = ['properties', ...schemaPath];
+
+        const jsonPointerPath = '/' + jsonPointerParts.join('/');
+
+        try {
+            const propertyDefinitions = resourceSchema.resolveJsonPointerPath(jsonPointerPath);
+
+            if (propertyDefinitions.length === 0) {
+                return undefined;
+            }
+
+            const propertyDef = propertyDefinitions[0];
+
+            if (propertyDef && 'type' in propertyDef) {
+                const type = propertyDef.type;
+                if (Array.isArray(type)) {
+                    return type[0];
+                } else if (typeof type === 'string') {
+                    return type;
+                }
+            }
+
+            return undefined;
+        } catch {
+            return undefined;
+        }
     }
 
     private formatForYaml(label: string, item: CompletionItem | undefined, editorSettings: EditorSettings): string {

src/autocomplete/CompletionRouter.ts

@@ -12,6 +12,7 @@ import { Entity, Output, Parameter } from '../context/semantic/Entity';
 import { EntityType } from '../context/semantic/SemanticTypes';
 import { DocumentType } from '../document/Document';
 import { DocumentManager } from '../document/DocumentManager';
+import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { CfnExternal } from '../server/CfnExternal';
 import { CfnInfraCore } from '../server/CfnInfraCore';
 import { CfnLspProviders } from '../server/CfnLspProviders';
@@ -48,6 +49,7 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
         private readonly contextManager: ContextManager,
         private readonly completionProviderMap: Map<CompletionProviderType, CompletionProvider>,
         private readonly documentManager: DocumentManager,
+        private readonly schemaRetriever: SchemaRetriever,
         private readonly entityFieldCompletionProviderMap = createEntityFieldProviders(),
     ) {}
 
@@ -89,6 +91,7 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
 
         const completions = provider?.getCompletions(context, params) ?? [];
         const editorSettings = this.documentManager.getEditorSettingsForDocument(params.textDocument.uri);
+        const lineContent = this.documentManager.getLine(params.textDocument.uri, context.startPosition.row);
 
         if (completions instanceof Promise) {
             return await completions.then((result) => {
@@ -99,6 +102,8 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
                     },
                     context,
                     editorSettings,
+                    lineContent,
+                    this.schemaRetriever,
                 );
             });
         } else if (completions) {
@@ -107,7 +112,7 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
                 items: completions.slice(0, this.completionSettings.maxCompletions),
             };
 
-            return this.formatter.format(completionList, context, editorSettings);
+            return this.formatter.format(completionList, context, editorSettings, lineContent, this.schemaRetriever);
         }
         return;
     }
@@ -272,10 +277,13 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
     }
 
     static create(core: CfnInfraCore, external: CfnExternal, providers: CfnLspProviders) {
+        CompletionFormatter.getInstance();
         return new CompletionRouter(
             core.contextManager,
             createCompletionProviders(core, external, providers),
             core.documentManager,
+            external.schemaRetriever,
+            createEntityFieldProviders(),
         );
     }
 }

src/autocomplete/CompletionUtils.ts

@@ -16,6 +16,18 @@ import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { ExtensionName } from '../utils/ExtensionConfig';
 import { ExtendedCompletionItem } from './CompletionFormatter';
 
+// Resource attributes are not given a data.type upon completionItem creation and
+// no schema for resource attributes. This is used for formatting completion items.
+export const RESOURCE_ATTRIBUTE_TYPES: Record<string, string> = {
+    DependsOn: 'string', // string | string[] - array of resource logical IDs
+    Condition: 'string', // Reference to a condition name
+    Metadata: 'object', // Arbitrary JSON/YAML object
+    CreationPolicy: 'object', // Configuration object with specific properties
+    DeletionPolicy: 'string', // Enum: "Delete" | "Retain" | "Snapshot"
+    UpdatePolicy: 'object', // Configuration object with update policies
+    UpdateReplacePolicy: 'string', // Enum: "Delete" | "Retain" | "Snapshot"
+};
+
 /**
  * Creates a replacement range from a context's start and end positions.
  * This is used for text edits in completion items.

src/autocomplete/TopLevelSectionCompletionProvider.ts

@@ -1,6 +1,6 @@
 import { CompletionItem, CompletionItemKind, CompletionParams, InsertTextFormat } from 'vscode-languageserver';
 import { Context } from '../context/Context';
-import { TopLevelSection, TopLevelSections } from '../context/ContextType';
+import { TopLevelSection, TopLevelSections, TopLevelSectionsWithLogicalIdsSet } from '../context/ContextType';
 import { SyntaxTreeManager } from '../context/syntaxtree/SyntaxTreeManager';
 import { DocumentType } from '../document/Document';
 import { DocumentManager } from '../document/DocumentManager';
@@ -119,7 +119,17 @@ ${CompletionFormatter.getIndentPlaceholder(1)}\${1:ConditionName}: $2`,
                 return this.constantsFeatureFlag.isEnabled();
             }
             return true;
-        }).map((section) => createCompletionItem(section, CompletionItemKind.Class));
+        }).map((section) => {
+            const shouldBeObject = TopLevelSectionsWithLogicalIdsSet.has(section);
+
+            const options = shouldBeObject
+                ? {
+                      data: { type: 'object' },
+                  }
+                : undefined;
+
+            return createCompletionItem(section, CompletionItemKind.Class, options);
+        });
     }
 
     private getTopLevelSectionSnippetCompletions(context: Context, params: CompletionParams): CompletionItem[] {