Report a specific diagnostic for walrus in a comprehension iterable expression

packages/pyright-internal/src/localization/localize.ts

@@ -1235,6 +1235,7 @@ export namespace Localizer {
             new ParameterizedString<{ names: string }>(getRawString('Diagnostic.variadicTypeParamTooManyClass'));
         export const walrusIllegal = () => getRawString('Diagnostic.walrusIllegal');
         export const walrusNotAllowed = () => getRawString('Diagnostic.walrusNotAllowed');
+        export const walrusNotAllowedInComprehension = () => getRawString('Diagnostic.walrusNotAllowedInComprehension');
         export const wildcardInFunction = () => getRawString('Diagnostic.wildcardInFunction');
         export const wildcardPatternTypeUnknown = () => getRawString('Diagnostic.wildcardPatternTypeUnknown');
         export const wildcardPatternTypePartiallyUnknown = () =>

packages/pyright-internal/src/localization/package.nls.en-us.json

@@ -1764,6 +1764,7 @@
         },
         "walrusIllegal": "Operator \":=\" requires Python 3.8 or newer",
         "walrusNotAllowed": "Operator \":=\" is not allowed in this context without surrounding parentheses",
+        "walrusNotAllowedInComprehension": "Operator \":=\" is not allowed within a comprehension iterable expression",
         "wildcardInFunction": {
             "message": "Wildcard import not allowed within a class or function",
             "comment": "{Locked='import'}"

packages/pyright-internal/src/parser/parser.ts

@@ -3380,7 +3380,12 @@ export class Parser {
             return leftExpr;
         }
 
-        if (!this._assignmentExpressionsAllowed || disallowAssignmentExpression) {
+        if (!this._assignmentExpressionsAllowed) {
+            // Assignment expressions are disallowed anywhere within the iterable
+            // expression of a comprehension's "for" clause, even if parenthesized.
+            // This differs from the "requires surrounding parentheses" case below.
+            this._addSyntaxError(LocMessage.walrusNotAllowedInComprehension(), walrusToken);
+        } else if (disallowAssignmentExpression) {
             this._addSyntaxError(LocMessage.walrusNotAllowed(), walrusToken);
         }
 

packages/pyright-internal/src/tests/samples/assignmentExprMessage1.py

@@ -0,0 +1,18 @@
+# This sample tests that assignment expressions used within the iterable
+# expression of a comprehension "for" clause are reported with a message
+# that specifically calls out the comprehension-iterable restriction from
+# PEP 572. Unlike a bare assignment expression used as a comprehension
+# "if" condition, this restriction cannot be resolved by adding parentheses.
+
+x = []
+
+
+# This should generate an error because an assignment expression is not
+# allowed within a comprehension's iterable expression, even when it is
+# surrounded by parentheses.
+[a for a in (b := x)]
+
+# This should generate an error because a bare (unparenthesized) assignment
+# expression is not allowed as a comprehension "if" condition. Here, adding
+# parentheses would make the code legal.
+[a for a in x if c := a]

packages/pyright-internal/src/tests/typeEvaluator1.test.ts

@@ -1215,3 +1215,22 @@ test('AssignmentExpr9', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['assignmentExpr9.py']);
     TestUtils.validateResults(analysisResults, 0);
 });
+
+test('AssignmentExprMessage1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['assignmentExprMessage1.py']);
+
+    TestUtils.validateResults(analysisResults, 2);
+
+    // A walrus within a comprehension's iterable expression cannot be fixed by
+    // adding parentheses, so it must use the comprehension-specific message
+    // rather than the generic "requires surrounding parentheses" message.
+    expect(analysisResults[0].errors[0].message).toBe(
+        'Operator ":=" is not allowed within a comprehension iterable expression'
+    );
+
+    // A bare walrus used as a comprehension "if" condition still uses the
+    // generic message because parenthesizing it makes the code legal.
+    expect(analysisResults[0].errors[1].message).toBe(
+        'Operator ":=" is not allowed in this context without surrounding parentheses'
+    );
+});