Eliminated the limitation that await, assignment expressions, f-str…

…ings, string chains, and escaped strings cannot be used within `Annotated` expressions when using an alias of `Annotated`. This addresses #6714. (#6719)
microsoft · Dec 13, 2023 · a5fa08e · a5fa08e
1 parent 23dea60
commit a5fa08e
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 7 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -1139,6 +1139,10 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             }
 
             case ParseNodeType.AssignmentExpression: {
+                if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) !== 0) {
+                    addError(Localizer.Diagnostic.walrusNotAllowed(), node);
+                }
+
                 typeResult = getTypeOfExpression(node.rightExpression, flags, inferenceContext);
                 assignTypeToExpression(
                     node.name,
@@ -1336,6 +1340,11 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
     }
 
     function getTypeOfAwaitOperator(node: AwaitNode, flags: EvaluatorFlags, inferenceContext?: InferenceContext) {
+        if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) !== 0) {
+            addError(Localizer.Diagnostic.awaitNotAllowed(), node);
+            return { type: UnknownType.create() };
+        }
+
         const effectiveExpectedType = inferenceContext
             ? createAwaitableReturnType(
                   node,

diff --git a/packages/pyright-internal/src/localization/localize.ts b/packages/pyright-internal/src/localization/localize.ts
@@ -246,6 +246,7 @@ export namespace Localizer {
         export const assignmentTargetExpr = () => getRawString('Diagnostic.assignmentTargetExpr');
         export const asyncNotInAsyncFunction = () => getRawString('Diagnostic.asyncNotInAsyncFunction');
         export const awaitIllegal = () => getRawString('Diagnostic.awaitIllegal');
+        export const awaitNotAllowed = () => getRawString('Diagnostic.awaitNotAllowed');
         export const awaitNotInAsync = () => getRawString('Diagnostic.awaitNotInAsync');
         export const backticksIllegal = () => getRawString('Diagnostic.backticksIllegal');
         export const baseClassCircular = () => getRawString('Diagnostic.baseClassCircular');

diff --git a/packages/pyright-internal/src/localization/package.nls.en-us.json b/packages/pyright-internal/src/localization/package.nls.en-us.json
@@ -30,6 +30,7 @@
         "assignmentTargetExpr": "Expression cannot be assignment target",
         "asyncNotInAsyncFunction": "Use of \"async\" not allowed outside of async function",
         "awaitIllegal": "Use of \"await\" requires Python 3.5 or newer",
+        "awaitNotAllowed": "Type annotations cannot use \"await\"",
         "awaitNotInAsync": "\"await\" allowed only within async function",
         "backticksIllegal": "Expressions surrounded by backticks are not supported in Python 3.x; use repr instead",
         "baseClassCircular": "Class cannot derive from itself",

diff --git a/packages/pyright-internal/src/parser/parser.ts b/packages/pyright-internal/src/parser/parser.ts
@@ -3164,7 +3164,7 @@ export class Parser {
             return leftExpr;
         }
 
-        if (!this._assignmentExpressionsAllowed || this._isParsingTypeAnnotation || disallowAssignmentExpression) {
+        if (!this._assignmentExpressionsAllowed || disallowAssignmentExpression) {
             this._addError(Localizer.Diagnostic.walrusNotAllowed(), walrusToken);
         }
 
@@ -3459,7 +3459,7 @@ export class Parser {
     // trailer: '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
     private _parseAtomExpression(): ExpressionNode {
         let awaitToken: KeywordToken | undefined;
-        if (this._peekKeywordType() === KeywordType.Await && !this._isParsingTypeAnnotation) {
+        if (this._peekKeywordType() === KeywordType.Await) {
             awaitToken = this._getKeywordToken(KeywordType.Await);
             if (this._getLanguageVersion() < PythonVersion.V3_5) {
                 this._addError(Localizer.Diagnostic.awaitIllegal(), awaitToken);
@@ -4875,9 +4875,13 @@ export class Parser {
             // Don't allow multiple strings because we have no way of reporting
             // parse errors that span strings.
             if (stringNode.strings.length > 1) {
-                this._addError(Localizer.Diagnostic.annotationSpansStrings(), stringNode);
+                if (this._isParsingQuotedText) {
+                    this._addError(Localizer.Diagnostic.annotationSpansStrings(), stringNode);
+                }
             } else if (stringNode.strings[0].nodeType === ParseNodeType.FormatString) {
-                this._addError(Localizer.Diagnostic.annotationFormatString(), stringNode);
+                if (this._isParsingQuotedText) {
+                    this._addError(Localizer.Diagnostic.annotationFormatString(), stringNode);
+                }
             } else {
                 const stringToken = stringNode.strings[0].token;
                 const stringValue = StringTokenUtils.getUnescapedString(stringNode.strings[0].token);
@@ -4888,7 +4892,9 @@ export class Parser {
                 // Don't allow escape characters because we have no way of mapping
                 // error ranges back to the escaped text.
                 if (unescapedString.length !== stringToken.length - prefixLength - stringToken.quoteMarkLength) {
-                    this._addError(Localizer.Diagnostic.annotationStringEscape(), stringNode);
+                    if (this._isParsingQuotedText) {
+                        this._addError(Localizer.Diagnostic.annotationStringEscape(), stringNode);
+                    }
                 } else {
                     const parser = new Parser();
                     const parseResults = parser.parseTextExpression(

diff --git a/packages/pyright-internal/src/tests/samples/annotated1.py b/packages/pyright-internal/src/tests/samples/annotated1.py
@@ -78,3 +78,13 @@ class B:
 reveal_type(Alias3, expected_text="type[str]")
 
 x2: Annotated[str, [*(1, 2)]]
+x3: Annotated[str, (temp := 1)]
+
+
+async def func3():
+    x4: Annotated[str, await func3()]
+
+
+x5: Annotated[str, f""]
+x6: Annotated[str, "a" "b" "c"]
+x7: Annotated[str, "a\nb"]
diff --git a/packages/pyright-internal/src/tests/samples/assignmentExpr3.py b/packages/pyright-internal/src/tests/samples/assignmentExpr3.py
@@ -16,7 +16,7 @@ def foo3(answer=(p := 42)):  # Valid, though not great style
 
 default_value: int = 3
 
-# This should generate an error.
+# This should generate two errors.
 def foo4(answer: p := default_value = 5):  # INVALID
     ...
 

diff --git a/packages/pyright-internal/src/tests/typeEvaluator4.test.ts b/packages/pyright-internal/src/tests/typeEvaluator4.test.ts
@@ -109,7 +109,7 @@ test('AssignmentExpr2', () => {
 
 test('AssignmentExpr3', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['assignmentExpr3.py']);
-    TestUtils.validateResults(analysisResults, 4);
+    TestUtils.validateResults(analysisResults, 5);
 });
 
 test('AssignmentExpr4', () => {