Skip to content

Commit

Permalink
Added support for other functional Enum calls including a list of str…
Browse files Browse the repository at this point in the history 
…ings, a tuple of strings, a list of tuples, a tuple of tuples, and a map. This addresses #1127.
  • Loading branch information
@msfterictraut
msfterictraut committed Apr 15, 2023
1 parent fcdb119 commit 1ee244c
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 13 deletions.
130 changes: 121 additions & 9 deletions packages/pyright-internal/src/analyzer/enums.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ import {
isFunction,
isInstantiableClass,
Type,
UnknownType,
} from './types';
import { computeMroLinearization } from './typeUtils';

Expand Down Expand Up @@ -85,7 +84,13 @@ export function createEnumType(
return undefined;
}

// The Enum factory call supports various forms of arguments:
const intClassType = evaluator.getBuiltInType(errorNode, 'int');
if (!intClassType || !isInstantiableClass(intClassType)) {
return undefined;
}
const classInstanceType = ClassType.cloneAsInstance(classType);

// The Enum functional form supports various forms of arguments:
// Enum('name', 'a b c')
// Enum('name', 'a,b,c')
// Enum('name', ['a', 'b', 'c'])
Expand All @@ -99,9 +104,6 @@ export function createEnumType(
return undefined;
}

const classInstanceType = ClassType.cloneAsInstance(classType);
const intClassType = evaluator.getBuiltInType(errorNode, 'int');

const initStr = initArg.valueExpression.strings
.map((s) => s.value)
.join('')
Expand All @@ -115,27 +117,137 @@ export function createEnumType(
return undefined;
}

const valueType =
intClassType && isInstantiableClass(intClassType)
? ClassType.cloneWithLiteral(ClassType.cloneAsInstance(intClassType), index + 1)
: UnknownType.create();
const valueType = ClassType.cloneWithLiteral(ClassType.cloneAsInstance(intClassType), index + 1);

const enumLiteral = new EnumLiteral(
classType.details.fullName,
classType.details.name,
entryName,
valueType
);

const newSymbol = Symbol.createWithType(
SymbolFlags.ClassMember,
ClassType.cloneWithLiteral(classInstanceType, enumLiteral)
);

classFields.set(entryName, newSymbol);
}

return classType;
}

if (
initArg.valueExpression.nodeType === ParseNodeType.List ||
initArg.valueExpression.nodeType === ParseNodeType.Tuple
) {
const entries =
initArg.valueExpression.nodeType === ParseNodeType.List
? initArg.valueExpression.entries
: initArg.valueExpression.expressions;

if (entries.length === 0) {
return undefined;
}

// Entries can be either string literals or tuples of a string
// literal and a value. All entries must follow the same pattern.
let isSimpleString = false;
for (const [index, entry] of entries.entries()) {
if (index === 0) {
isSimpleString = entry.nodeType === ParseNodeType.StringList;
}

let nameNode: ParseNode | undefined;
let valueType: Type | undefined;

if (entry.nodeType === ParseNodeType.StringList) {
if (!isSimpleString) {
return undefined;
}

nameNode = entry;
valueType = ClassType.cloneWithLiteral(ClassType.cloneAsInstance(intClassType), index + 1);
} else if (entry.nodeType === ParseNodeType.Tuple) {
if (isSimpleString) {
return undefined;
}

if (entry.expressions.length !== 2) {
return undefined;
}
nameNode = entry.expressions[0];
valueType = evaluator.getTypeOfExpression(entry.expressions[1]).type;
} else {
return undefined;
}

if (
nameNode.nodeType !== ParseNodeType.StringList ||
nameNode.strings.length !== 1 ||
nameNode.strings[0].nodeType !== ParseNodeType.String
) {
return undefined;
}

const entryName = nameNode.strings[0].value;

const enumLiteral = new EnumLiteral(
classType.details.fullName,
classType.details.name,
entryName,
valueType
);

const newSymbol = Symbol.createWithType(
SymbolFlags.ClassMember,
ClassType.cloneWithLiteral(classInstanceType, enumLiteral)
);

classFields.set(entryName, newSymbol);
}
}

if (initArg.valueExpression.nodeType === ParseNodeType.Dictionary) {
const entries = initArg.valueExpression.entries;
if (entries.length === 0) {
return undefined;
}

for (const entry of entries) {
// Don't support dictionary expansion expressions.
if (entry.nodeType !== ParseNodeType.DictionaryKeyEntry) {
return undefined;
}

const nameNode = entry.keyExpression;
const valueType = evaluator.getTypeOfExpression(entry.valueExpression).type;

if (
nameNode.nodeType !== ParseNodeType.StringList ||
nameNode.strings.length !== 1 ||
nameNode.strings[0].nodeType !== ParseNodeType.String
) {
return undefined;
}

const entryName = nameNode.strings[0].value;
const enumLiteral = new EnumLiteral(
classType.details.fullName,
classType.details.name,
entryName,
valueType
);

const newSymbol = Symbol.createWithType(
SymbolFlags.ClassMember,
ClassType.cloneWithLiteral(classInstanceType, enumLiteral)
);

classFields.set(entryName, newSymbol);
}
}

return classType;
}

Expand Down
54 changes: 50 additions & 4 deletions packages/pyright-internal/src/tests/samples/enums1.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from enum import Enum, IntEnum


TestEnum1 = Enum("TestEnum1", " A B,C , \t D\t")
TestEnum1 = Enum("TestEnum1", " A B, , ,C , \t D\t")
TestEnum2 = IntEnum("TestEnum2", "AA BB CC DD")


Expand Down Expand Up @@ -55,6 +55,52 @@ class TestEnum3(Enum):
reveal_type(num_items_in_enum3, expected_text="int")

reveal_type(TestEnum3.A.name, expected_text="Literal['A']")
reveal_type(TestEnum3.A._name_, expected_text="Literal['A']")
reveal_type(TestEnum3.A.value, expected_text="Literal[0]")
reveal_type(TestEnum3.A._value_, expected_text="Literal[0]")
reveal_type(TestEnum3.B._name_, expected_text="Literal['B']")
reveal_type(TestEnum3.C.value, expected_text="Literal[2]")
reveal_type(TestEnum3.D._value_, expected_text="Literal[3]")


TestEnum4 = Enum("TestEnum4", ["A", "B", "C", "D"])
reveal_type(TestEnum4.A, expected_text="Literal[TestEnum4.A]")
reveal_type(TestEnum4.D, expected_text="Literal[TestEnum4.D]")
reveal_type(TestEnum4.A.name, expected_text="Literal['A']")
reveal_type(TestEnum4.B._name_, expected_text="Literal['B']")
reveal_type(TestEnum4.C.value, expected_text="Literal[3]")
reveal_type(TestEnum4.D._value_, expected_text="Literal[4]")

TestEnum5 = Enum("TestEnum5", ("A", "B", "C", "D"))
reveal_type(TestEnum5.A, expected_text="Literal[TestEnum5.A]")
reveal_type(TestEnum5.D, expected_text="Literal[TestEnum5.D]")
reveal_type(TestEnum5.A.name, expected_text="Literal['A']")
reveal_type(TestEnum5.B._name_, expected_text="Literal['B']")
reveal_type(TestEnum5.C.value, expected_text="Literal[3]")
reveal_type(TestEnum5.D._value_, expected_text="Literal[4]")

d_value = "d"

TestEnum6 = Enum("TestEnum6", [("A", 1), ("B", [1, 2]), ("C", "c"), ("D", d_value)])
reveal_type(TestEnum6.A, expected_text="Literal[TestEnum6.A]")
reveal_type(TestEnum6.D, expected_text="Literal[TestEnum6.D]")
reveal_type(TestEnum6.A.name, expected_text="Literal['A']")
reveal_type(TestEnum6.B._name_, expected_text="Literal['B']")
reveal_type(TestEnum6.A.value, expected_text="Literal[1]")
reveal_type(TestEnum6.B.value, expected_text="list[int]")
reveal_type(TestEnum6.C.value, expected_text="Literal['c']")
reveal_type(TestEnum6.D._value_, expected_text="Literal['d']")

TestEnum7 = Enum("TestEnum7", (("A", 1), ("D", "d")))
reveal_type(TestEnum7.A, expected_text="Literal[TestEnum7.A]")
reveal_type(TestEnum7.D, expected_text="Literal[TestEnum7.D]")
reveal_type(TestEnum7.A.name, expected_text="Literal['A']")
reveal_type(TestEnum7.A.value, expected_text="Literal[1]")
reveal_type(TestEnum7.D._value_, expected_text="Literal['d']")

TestEnum8 = Enum("TestEnum8", {"A": 1, "B": [1, 2], "C": "c", "D": d_value})
reveal_type(TestEnum8.A, expected_text="Literal[TestEnum8.A]")
reveal_type(TestEnum8.D, expected_text="Literal[TestEnum8.D]")
reveal_type(TestEnum8.A.name, expected_text="Literal['A']")
reveal_type(TestEnum8.B._name_, expected_text="Literal['B']")
reveal_type(TestEnum8.A.value, expected_text="Literal[1]")
reveal_type(TestEnum8.B.value, expected_text="list[int]")
reveal_type(TestEnum8.C.value, expected_text="Literal['c']")
reveal_type(TestEnum8.D._value_, expected_text="Literal['d']")

0 comments on commit 1ee244c

Please sign in to comment.