From eb45e303a429f6793caa26e9c7764849ab3779da Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Tue, 5 Dec 2023 19:15:21 -0800
Subject: [PATCH] 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.

---
 .../src/analyzer/typePrinter.ts               | 57 ++++++++++++++++---
 .../src/analyzer/typeStubWriter.ts            |  2 +-
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/packages/pyright-internal/src/analyzer/typePrinter.ts b/packages/pyright-internal/src/analyzer/typePrinter.ts
index 5579cdc75b6a..af5309ffcb18 100644
--- 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
index eb64c4a17670..1bbeee9a779f 100644
--- 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 })}:`;
                     }
                 }
             }