Changed the type printer logic to avoid emitting truncated Literal …

…string and bytes values when it is asked to enforce Python syntax. This mode is used for inlay hints and the stub generator (for inferred return type comments). This addresses #6659.
microsoft · Dec 6, 2023 · eb45e30 · eb45e30
1 parent 5cd59ad
commit eb45e30
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 8 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typePrinter.ts b/packages/pyright-internal/src/analyzer/typePrinter.ts
@@ -8,6 +8,7 @@
  */
 
 import { appendArray } from '../common/collectionUtils';
+import { assert } from '../common/debug';
 import { ParameterCategory } from '../parser/parseNodes';
 import { isTypedKwargs } from './parameterUtils';
 import * as ParseTreeUtils from './parseTreeUtils';
@@ -122,6 +123,27 @@ export function printObjectTypeForClass(
     return printObjectTypeForClassInternal(type, printTypeFlags, returnTypeCallback, uniqueNameMap, [], 0);
 }
 
+const maxLiteralStringLength = 50;
+
+export function isLiteralValueTruncated(type: ClassType): boolean {
+    if (typeof type.literalValue === 'string') {
+        if (type.literalValue.length > maxLiteralStringLength) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+export function printLiteralValueTruncated(type: ClassType): string {
+    if (type.details.name === 'bytes') {
+        return 'bytes';
+    }
+
+    assert(type.details.name === 'str');
+    return 'LiteralString';
+}
+
 export function printLiteralValue(type: ClassType, quotation = "'"): string {
     const literalValue = type.literalValue;
     if (literalValue === undefined) {
@@ -133,7 +155,6 @@ export function printLiteralValue(type: ClassType, quotation = "'"): string {
         let effectiveLiteralValue = literalValue;
 
         // Limit the length of the string literal.
-        const maxLiteralStringLength = 50;
         if (literalValue.length > maxLiteralStringLength) {
             effectiveLiteralValue = literalValue.substring(0, maxLiteralStringLength) + '…';
         }
@@ -402,7 +423,11 @@ function printTypeInternal(
             case TypeCategory.Class: {
                 if (TypeBase.isInstance(type)) {
                     if (type.literalValue !== undefined) {
-                        return `Literal[${printLiteralValue(type)}]`;
+                        if (isLiteralValueTruncated(type) && (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0) {
+                            return printLiteralValueTruncated(type);
+                        } else {
+                            return `Literal[${printLiteralValue(type)}]`;
+                        }
                     }
 
                     return `${printObjectTypeForClassInternal(
@@ -417,16 +442,20 @@ function printTypeInternal(
                     let typeToWrap: string;
 
                     if (type.literalValue !== undefined) {
-                        typeToWrap = `Literal[${printLiteralValue(type)}]`;
+                        if (isLiteralValueTruncated(type) && (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0) {
+                            typeToWrap = printLiteralValueTruncated(type);
+                        } else {
+                            typeToWrap = `Literal[${printLiteralValue(type)}]`;
+                        }
                     } else {
-                        typeToWrap = `${printObjectTypeForClassInternal(
+                        typeToWrap = printObjectTypeForClassInternal(
                             type,
                             printTypeFlags,
                             returnTypeCallback,
                             uniqueNameMap,
                             recursionTypes,
                             recursionCount
-                        )}`;
+                        );
                     }
 
                     return `${_printNestedInstantiable(type, typeToWrap)}${getConditionalIndicator(type)}`;
@@ -573,9 +602,23 @@ function printTypeInternal(
                 doForEachSubtype(type, (subtype, index) => {
                     if (!subtypeHandledSet.has(index)) {
                         if (isClassInstance(subtype) && subtype.literalValue !== undefined) {
-                            literalObjectStrings.add(printLiteralValue(subtype));
+                            if (
+                                isLiteralValueTruncated(subtype) &&
+                                (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0
+                            ) {
+                                subtypeStrings.add(printLiteralValueTruncated(subtype));
+                            } else {
+                                literalObjectStrings.add(printLiteralValue(subtype));
+                            }
                         } else if (isInstantiableClass(subtype) && subtype.literalValue !== undefined) {
-                            literalClassStrings.add(printLiteralValue(subtype));
+                            if (
+                                isLiteralValueTruncated(subtype) &&
+                                (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0
+                            ) {
+                                subtypeStrings.add(`type[${printLiteralValueTruncated(subtype)}]`);
+                            } else {
+                                literalClassStrings.add(printLiteralValue(subtype));
+                            }
                         } else {
                             subtypeStrings.add(
                                 printTypeInternal(

diff --git a/packages/pyright-internal/src/analyzer/typeStubWriter.ts b/packages/pyright-internal/src/analyzer/typeStubWriter.ts
@@ -289,7 +289,7 @@ export class TypeStubWriter extends ParseTreeWalker {
                     let returnType = this._evaluator.getFunctionInferredReturnType(functionType.functionType);
                     returnType = removeUnknownFromUnion(returnType);
                     if (!isNever(returnType) && !isUnknown(returnType)) {
-                        line += ` # -> ${this._evaluator.printType(returnType)}:`;
+                        line += ` # -> ${this._evaluator.printType(returnType, { enforcePythonSyntax: true })}:`;
                     }
                 }
             }