diff --git a/packages/pyright-internal/src/localization/localize.ts b/packages/pyright-internal/src/localization/localize.ts index 34fb7fb15932..fb16ccb333f3 100644 --- a/packages/pyright-internal/src/localization/localize.ts +++ b/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 = () => diff --git a/packages/pyright-internal/src/localization/package.nls.en-us.json b/packages/pyright-internal/src/localization/package.nls.en-us.json index 69d6957c61c9..e3075a10ccf7 100644 --- a/packages/pyright-internal/src/localization/package.nls.en-us.json +++ b/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'}" diff --git a/packages/pyright-internal/src/parser/parser.ts b/packages/pyright-internal/src/parser/parser.ts index ee6b04504644..481697724e65 100644 --- a/packages/pyright-internal/src/parser/parser.ts +++ b/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); } diff --git a/packages/pyright-internal/src/tests/samples/assignmentExprMessage1.py b/packages/pyright-internal/src/tests/samples/assignmentExprMessage1.py new file mode 100644 index 000000000000..ff0e05881f03 --- /dev/null +++ b/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] diff --git a/packages/pyright-internal/src/tests/typeEvaluator1.test.ts b/packages/pyright-internal/src/tests/typeEvaluator1.test.ts index 73306b3b11a4..ad5d65969065 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator1.test.ts +++ b/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' + ); +});