Added special-case support for enum fields that are assigned using an…

… unpacked tuple operation. This addresses #5458.
microsoft · Jul 11, 2023 · f0eec26 · f0eec26
1 parent c41bc9c
commit f0eec26
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 10 deletions.
diff --git a/packages/pyright-internal/src/analyzer/enums.ts b/packages/pyright-internal/src/analyzer/enums.ts
@@ -8,10 +8,23 @@
  */
 
 import { assert } from '../common/debug';
-import { ArgumentCategory, ExpressionNode, NameNode, ParseNode, ParseNodeType } from '../parser/parseNodes';
+import {
+    ArgumentCategory,
+    AssignmentNode,
+    ExpressionNode,
+    NameNode,
+    ParseNode,
+    ParseNodeType,
+} from '../parser/parseNodes';
 import { getFileInfo } from './analyzerNodeInfo';
 import { VariableDeclaration } from './declaration';
-import { getClassFullName, getEnclosingClass, getTypeSourceId } from './parseTreeUtils';
+import {
+    getClassFullName,
+    getEnclosingClass,
+    getParentNodeOfType,
+    getTypeSourceId,
+    isNodeContainedWithin,
+} from './parseTreeUtils';
 import { Symbol, SymbolFlags } from './symbol';
 import { isSingleDunderName } from './symbolNameUtils';
 import { FunctionArgument, TypeEvaluator } from './typeEvaluatorTypes';
@@ -23,6 +36,7 @@ import {
     ClassTypeFlags,
     EnumLiteral,
     Type,
+    UnknownType,
     combineTypes,
     isClass,
     isClassInstance,
@@ -279,14 +293,24 @@ export function transformTypeForPossibleEnumClass(
     // variables used within each enum instance. Unless/until there is
     // a change to this convention and all type checkers and stubs adopt
     // it, we're stuck with this limitation.
-    let isMemberOfEnumeration =
-        (node.parent?.nodeType === ParseNodeType.Assignment && node.parent.leftExpression === node) ||
-        (node.parent?.nodeType === ParseNodeType.TypeAnnotation &&
-            node.parent.valueExpression === node &&
-            node.parent.parent?.nodeType === ParseNodeType.Assignment) ||
-        (getFileInfo(node).isStubFile &&
-            node.parent?.nodeType === ParseNodeType.TypeAnnotation &&
-            node.parent.valueExpression === node);
+    let isMemberOfEnumeration = false;
+    let isUnpackedTuple = false;
+
+    const assignmentNode = getParentNodeOfType(node, ParseNodeType.Assignment) as AssignmentNode | undefined;
+
+    if (assignmentNode && isNodeContainedWithin(node, assignmentNode.leftExpression)) {
+        isMemberOfEnumeration = true;
+
+        if (getParentNodeOfType(node, ParseNodeType.Tuple)) {
+            isUnpackedTuple = true;
+        }
+    } else if (
+        getFileInfo(node).isStubFile &&
+        node.parent?.nodeType === ParseNodeType.TypeAnnotation &&
+        node.parent.valueExpression === node
+    ) {
+        isMemberOfEnumeration = true;
+    }
 
     // The spec specifically excludes names that start and end with a single underscore.
     // This also includes dunder names.
@@ -313,6 +337,14 @@ export function transformTypeForPossibleEnumClass(
         valueType = AnyType.create();
     } else {
         valueType = getValueType();
+
+        // If the LHS is an unpacked tuple, we need to handle this as
+        // a special case.
+        if (isUnpackedTuple) {
+            valueType =
+                evaluator.getTypeOfIterator({ type: valueType }, /* isAsync */ false, /* errorNode */ undefined)
+                    ?.type ?? UnknownType.create();
+        }
     }
 
     // The spec excludes descriptors.

diff --git a/packages/pyright-internal/src/tests/samples/enum1.py b/packages/pyright-internal/src/tests/samples/enum1.py
@@ -163,3 +163,12 @@ class TestEnum10(Enum, metaclass=CustomEnumMeta1):
 def func2(e: type[Enum]):
     values = {v.value for v in e}
     reveal_type(values, expected_text="set[Any]")
+
+
+class TestEnum11(Enum):
+    (A, B, C) = range(3)
+
+
+te11_A = TestEnum11.A
+reveal_type(te11_A, expected_text="Literal[TestEnum11.A]")
+reveal_type(te11_A.value, expected_text="int")