Added check for concatenated bytes and str literals. This addresses m…

…icrosoft#8871.
lucentlabz · Sep 2, 2024 · 386b7fe · 386b7fe
1 parent 689857d
commit 386b7fe
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 3 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -1546,13 +1546,30 @@ export function createTypeEvaluator(
         return typeResult;
     }
 
-    function getTypeOfStringList(node: StringListNode, flags: EvalFlags) {
+    function getTypeOfStringList(node: StringListNode, flags: EvalFlags): TypeResult {
         let typeResult: TypeResult | undefined;
 
         if ((flags & EvalFlags.StrLiteralAsType) !== 0 && (flags & EvalFlags.TypeFormArg) === 0) {
             return getTypeOfStringListAsType(node, flags);
         }
 
+        const isBytesNode = (node: StringNode | FormatStringNode) =>
+            (node.d.token.flags & StringTokenFlags.Bytes) !== 0;
+
+        // Check for mixing of bytes and str, which is not allowed.
+        const firstStrIndex = node.d.strings.findIndex((str) => !isBytesNode(str));
+        const firstBytesIndex = node.d.strings.findIndex((str) => isBytesNode(str));
+        if (firstStrIndex >= 0 && firstBytesIndex >= 0) {
+            addDiagnostic(
+                DiagnosticRule.reportGeneralTypeIssues,
+                LocMessage.mixingBytesAndStr(),
+                node.d.strings[Math.max(firstBytesIndex, firstStrIndex)]
+            );
+
+            return { type: UnknownType.create() };
+        }
+
+        const isBytes = firstBytesIndex >= 0;
         let isLiteralString = true;
         let isIncomplete = false;
 
@@ -1577,8 +1594,6 @@ export function createTypeEvaluator(
             }
         });
 
-        const isBytes = (node.d.strings[0].d.token.flags & StringTokenFlags.Bytes) !== 0;
-
         // Don't create a literal type if it's an f-string.
         if (node.d.strings.some((str) => str.nodeType === ParseNodeType.FormatString)) {
             if (isLiteralString) {

diff --git a/packages/pyright-internal/src/localization/localize.ts b/packages/pyright-internal/src/localization/localize.ts
@@ -649,6 +649,7 @@ export namespace Localizer {
             new ParameterizedString<{ name: string }>(getRawString('Diagnostic.methodReturnsNonObject'));
         export const missingSuperCall = () =>
             new ParameterizedString<{ methodName: string }>(getRawString('Diagnostic.missingSuperCall'));
+        export const mixingBytesAndStr = () => getRawString('Diagnostic.mixingBytesAndStr');
         export const moduleAsType = () => getRawString('Diagnostic.moduleAsType');
         export const moduleNotCallable = () => getRawString('Diagnostic.moduleNotCallable');
         export const moduleUnknownMember = () =>

diff --git a/packages/pyright-internal/src/localization/package.nls.en-us.json b/packages/pyright-internal/src/localization/package.nls.en-us.json
@@ -316,6 +316,7 @@
         "methodOverridden": "\"{name}\" overrides method of same name in class \"{className}\" with incompatible type \"{type}\"",
         "methodReturnsNonObject": "\"{name}\" method does not return an object",
         "missingSuperCall": "Method \"{methodName}\" does not call the method of the same name in parent class",
+        "mixingBytesAndStr": "Bytes and str values cannot be concatenated",
         "moduleAsType": "Module cannot be used as a type",
         "moduleNotCallable": "Module is not callable",
         "moduleUnknownMember": "\"{memberName}\" is not a known attribute of module \"{moduleName}\"",

diff --git a/packages/pyright-internal/src/tests/samples/strings2.py b/packages/pyright-internal/src/tests/samples/strings2.py
@@ -0,0 +1,14 @@
+# This sample tests string concatenation.
+
+v1 = "a" "b" r"c" R"""d""" "e" "f"
+reveal_type(v1, expected_text="Literal['abcdef']")
+
+v2 = b"a" b"b" rb"c" Rb"d"
+reveal_type(v2, expected_text='Literal[b"abcd"]')
+
+# This should generate an error.
+v3 = "a" b"b"
+
+# This should generate an error.
+v4 = b"a" f""
+
diff --git a/packages/pyright-internal/src/tests/typeEvaluator8.test.ts b/packages/pyright-internal/src/tests/typeEvaluator8.test.ts
@@ -805,6 +805,12 @@ test('PseudoGeneric3', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Strings2', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['strings2.py']);
+
+    TestUtils.validateResults(analysisResults, 2);
+});
+
 test('LiteralString1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['literalString1.py']);