Eliminated the limitation that unpack operators cannot be used within…

… `Annotated` expressions when using an alias of `Annotated`. This partially addresses #6714. (#6718)
microsoft · Dec 13, 2023 · 23dea60 · 23dea60
1 parent 95f9f87
commit 23dea60
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 30 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -1402,26 +1402,28 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             !iterType.isVariadicUnpacked
         ) {
             typeResult = { type: TypeVarType.cloneForUnpacked(iterType) };
+        } else if (
+            (flags & EvaluatorFlags.AllowUnpackedTupleOrTypeVarTuple) !== 0 &&
+            isInstantiableClass(iterType) &&
+            ClassType.isBuiltIn(iterType, 'tuple')
+        ) {
+            typeResult = { type: ClassType.cloneForUnpacked(iterType) };
+        } else if ((flags & EvaluatorFlags.ExpectingTypeAnnotation) !== 0) {
+            addError(Localizer.Diagnostic.unpackInAnnotation(), node, node.starToken);
+            typeResult = { type: UnknownType.create() };
         } else {
-            if (
-                (flags & EvaluatorFlags.AllowUnpackedTupleOrTypeVarTuple) !== 0 &&
-                isInstantiableClass(iterType) &&
-                ClassType.isBuiltIn(iterType, 'tuple')
-            ) {
-                typeResult = { type: ClassType.cloneForUnpacked(iterType) };
-            } else {
-                const iteratorTypeResult = getTypeOfIterator(iterTypeResult, /* isAsync */ false, node) ?? {
-                    type: UnknownType.create(!!iterTypeResult.isIncomplete),
-                    isIncomplete: iterTypeResult.isIncomplete,
-                };
-                typeResult = {
-                    type: iteratorTypeResult.type,
-                    typeErrors: iterTypeResult.typeErrors,
-                    unpackedType: iterType,
-                    isIncomplete: iteratorTypeResult.isIncomplete,
-                };
-            }
+            const iteratorTypeResult = getTypeOfIterator(iterTypeResult, /* isAsync */ false, node) ?? {
+                type: UnknownType.create(!!iterTypeResult.isIncomplete),
+                isIncomplete: iterTypeResult.isIncomplete,
+            };
+            typeResult = {
+                type: iteratorTypeResult.type,
+                typeErrors: iterTypeResult.typeErrors,
+                unpackedType: iterType,
+                isIncomplete: iteratorTypeResult.isIncomplete,
+            };
         }
+
         return typeResult;
     }
 

diff --git a/packages/pyright-internal/src/parser/parseNodes.ts b/packages/pyright-internal/src/parser/parseNodes.ts
@@ -1189,11 +1189,13 @@ export namespace TernaryNode {
 export interface UnpackNode extends ParseNodeBase {
     readonly nodeType: ParseNodeType.Unpack;
     expression: ExpressionNode;
+    starToken: Token;
 }
 
 export namespace UnpackNode {
     export function create(starToken: Token, expression: ExpressionNode) {
         const node: UnpackNode = {
+            starToken,
             start: starToken.start,
             length: starToken.length,
             nodeType: ParseNodeType.Unpack,

diff --git a/packages/pyright-internal/src/parser/parser.ts b/packages/pyright-internal/src/parser/parser.ts
@@ -4503,20 +4503,18 @@ export class Parser {
         const startToken = this._peekToken();
         const isUnpack = this._consumeTokenIfOperator(OperatorType.Multiply);
 
-        if (isUnpack) {
-            if (!allowUnpack) {
-                this._addError(Localizer.Diagnostic.unpackInAnnotation(), startToken);
-            } else if (
-                !this._parseOptions.isStubFile &&
-                !this._isParsingQuotedText &&
-                this._getLanguageVersion() < PythonVersion.V3_11
-            ) {
-                this._addError(Localizer.Diagnostic.unpackedSubscriptIllegal(), startToken);
-            }
+        if (
+            isUnpack &&
+            allowUnpack &&
+            !this._parseOptions.isStubFile &&
+            !this._isParsingQuotedText &&
+            this._getLanguageVersion() < PythonVersion.V3_11
+        ) {
+            this._addError(Localizer.Diagnostic.unpackedSubscriptIllegal(), startToken);
         }
 
         let result = this._parseTestExpression(/* allowAssignmentExpression */ false);
-        if (isUnpack && allowUnpack) {
+        if (isUnpack) {
             result = UnpackNode.create(startToken, result);
         }
 

diff --git a/packages/pyright-internal/src/tests/samples/annotated1.py b/packages/pyright-internal/src/tests/samples/annotated1.py
@@ -53,7 +53,7 @@ def func2(a: TypeWithStringArg):
 _T = TypeVar("_T")
 Param = Annotated[_T, "x"]
 
-x: Param[int] = 3
+x1: Param[int] = 3
 
 
 class A:
@@ -76,3 +76,5 @@ class B:
 Alias3 = Alias1[Alias2]
 
 reveal_type(Alias3, expected_text="type[str]")
+
+x2: Annotated[str, [*(1, 2)]]