Skip to content

Add support for iterating over enums #42465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/compiler/checker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35078,6 +35078,17 @@ namespace ts {
return anyIterationTypes;
}

const symbolType = getDeclaredTypeOfSymbol(type.symbol);

if (symbolType.symbol.flags & SymbolFlags.RegularEnum) {
return createIterationTypes(
createTupleType([
stringType,
symbolType,
]),
);
}

if (use & IterationUse.AllowsAsyncIterablesFlag) {
const iterationTypes =
getIterationTypesOfIterableCached(type, asyncIterationTypesResolver) ||
Expand Down
90 changes: 85 additions & 5 deletions src/compiler/transformers/ts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2425,7 +2425,20 @@ namespace ts {

const statements: Statement[] = [];
startLexicalEnvironment();
const members = map(node.members, transformEnumMember);
const membersAndExpressions = node.members.map(member => {
const [nameExpression, valueExpression] = getEnumMemberExpressions(member);
return {
member,
nameExpression,
valueExpression,
};
});
const members = membersAndExpressions.map(m => transformEnumMember(
m.member,
m.nameExpression,
m.valueExpression,
));
members.push(transformEnumMembersToIterator(membersAndExpressions));
insertStatementsAfterStandardPrologue(statements, endLexicalEnvironment());
addRange(statements, members);

Expand All @@ -2437,20 +2450,37 @@ namespace ts {
}

/**
* Transforms an enum member into a statement.
* Get an expression for the name and value of an enum member.
* For example `enum Maybe { Yes = 1, No = 2 }` would give
* [`"Yes"`, `1`] and [`"No"`, `2`] as the expressions.
*
* @param member The enum member node.
*/
function transformEnumMember(member: EnumMember): Statement {
function getEnumMemberExpressions(member: EnumMember): [Expression, Expression] {
// enums don't support computed properties
// we pass false as 'generateNameForComputedPropertyName' for a backward compatibility purposes
// old emitter always generate 'expression' part of the name as-is.
const name = getExpressionForPropertyName(member, /*generateNameForComputedPropertyName*/ false);
const valueExpression = transformEnumMemberDeclarationValue(member);
return [name, valueExpression];
}

/**
* Transforms an enum member into a statement.
*
* @param member The enum member node.
* @param nameExpression A string literal expression containing the member name
* @param valueExpression An expression containing the member value.
*/
function transformEnumMember(
member: EnumMember,
nameExpression: Expression,
valueExpression: Expression,
): Statement {
const innerAssignment = factory.createAssignment(
factory.createElementAccessExpression(
currentNamespaceContainerName,
name
nameExpression
),
valueExpression
);
Expand All @@ -2461,7 +2491,7 @@ namespace ts {
currentNamespaceContainerName,
innerAssignment
),
name
nameExpression
);
return setTextRange(
factory.createExpressionStatement(
Expand All @@ -2474,6 +2504,56 @@ namespace ts {
);
}

/**
* Transforms all enum members into a generator function that yields
* the enum members.
*
* @param members The enum member nodes.
*/
function transformEnumMembersToIterator(
memberExpressions: {
nameExpression: Expression,
valueExpression: Expression,
}[],
): Statement {
const symbolIterator = factory.createPropertyAccessExpression(
factory.createIdentifier("Symbol"),
"iterator",
);
const iteratorMember = factory.createElementAccessExpression(
currentNamespaceContainerName,
symbolIterator,
);

const yieldExpression = factory.createYieldExpression(
factory.createToken(SyntaxKind.AsteriskToken),
factory.createArrayLiteralExpression(memberExpressions.map(
expressions => factory.createArrayLiteralExpression([
expressions.nameExpression,
expressions.valueExpression,
]),
)),
);
const generatorBody = factory.createBlock([
factory.createExpressionStatement(yieldExpression)
]);

const iteratorGenerator = factory.createFunctionExpression(
/*modifiers*/ undefined,
factory.createToken(SyntaxKind.AsteriskToken),
/*name*/ undefined,
[],
[],
/*type*/ undefined,
generatorBody,
);
const outerAssignment = factory.createAssignment(
iteratorMember,
iteratorGenerator,
);
return factory.createExpressionStatement(outerAssignment);
}

/**
* Transforms the value of an enum member.
*
Expand Down
35 changes: 35 additions & 0 deletions tests/baselines/reference/enumIteratorBasics1.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
//// [enumIteratorBasics1.ts]
enum Test {
Foo,
Bar,
Baz,
}

const AlsoTest = Test;

for (const member of AlsoTest) {
const x: string = member[0];
const y: Test = member[1];
console.log(x, y);
}

const members: [string, Test][] = [...AlsoTest];
console.log(members);


//// [enumIteratorBasics1.js]
var Test;
(function (Test) {
Test[Test["Foo"] = 0] = "Foo";
Test[Test["Bar"] = 1] = "Bar";
Test[Test["Baz"] = 2] = "Baz";
Test[Symbol.iterator] = function* () { yield* [["Foo", 0], ["Bar", 1], ["Baz", 2]]; };
})(Test || (Test = {}));
const AlsoTest = Test;
for (const member of AlsoTest) {
const x = member[0];
const y = member[1];
console.log(x, y);
}
const members = [...AlsoTest];
console.log(members);
52 changes: 52 additions & 0 deletions tests/baselines/reference/enumIteratorBasics1.symbols
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
=== tests/cases/compiler/enumIteratorBasics1.ts ===
enum Test {
>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))

Foo,
>Foo : Symbol(Test.Foo, Decl(enumIteratorBasics1.ts, 0, 11))

Bar,
>Bar : Symbol(Test.Bar, Decl(enumIteratorBasics1.ts, 1, 6))

Baz,
>Baz : Symbol(Test.Baz, Decl(enumIteratorBasics1.ts, 2, 6))
}

const AlsoTest = Test;
>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))
>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))

for (const member of AlsoTest) {
>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))

const x: string = member[0];
>x : Symbol(x, Decl(enumIteratorBasics1.ts, 9, 7))
>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
>0 : Symbol(0)

const y: Test = member[1];
>y : Symbol(y, Decl(enumIteratorBasics1.ts, 10, 7))
>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
>1 : Symbol(1)

console.log(x, y);
>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
>x : Symbol(x, Decl(enumIteratorBasics1.ts, 9, 7))
>y : Symbol(y, Decl(enumIteratorBasics1.ts, 10, 7))
}

const members: [string, Test][] = [...AlsoTest];
>members : Symbol(members, Decl(enumIteratorBasics1.ts, 14, 5))
>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))

console.log(members);
>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
>members : Symbol(members, Decl(enumIteratorBasics1.ts, 14, 5))

56 changes: 56 additions & 0 deletions tests/baselines/reference/enumIteratorBasics1.types
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
=== tests/cases/compiler/enumIteratorBasics1.ts ===
enum Test {
>Test : Test

Foo,
>Foo : Test.Foo

Bar,
>Bar : Test.Bar

Baz,
>Baz : Test.Baz
}

const AlsoTest = Test;
>AlsoTest : typeof Test
>Test : typeof Test

for (const member of AlsoTest) {
>member : [string, Test]
>AlsoTest : typeof Test

const x: string = member[0];
>x : string
>member[0] : string
>member : [string, Test]
>0 : 0

const y: Test = member[1];
>y : Test
>member[1] : Test
>member : [string, Test]
>1 : 1

console.log(x, y);
>console.log(x, y) : void
>console.log : (...data: any[]) => void
>console : Console
>log : (...data: any[]) => void
>x : string
>y : Test
}

const members: [string, Test][] = [...AlsoTest];
>members : [string, Test][]
>[...AlsoTest] : [string, Test][]
>...AlsoTest : [string, Test]
>AlsoTest : typeof Test

console.log(members);
>console.log(members) : void
>console.log : (...data: any[]) => void
>console : Console
>log : (...data: any[]) => void
>members : [string, Test][]

19 changes: 19 additions & 0 deletions tests/baselines/reference/enumIteratorConstEnum1.errors.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
tests/cases/compiler/enumIteratorConstEnum1.ts(8,22): error TS2475: 'const' enums can only be used in property or index access expressions or the right hand side of an import declaration or export assignment or type query.
tests/cases/compiler/enumIteratorConstEnum1.ts(8,22): error TS2488: Type 'typeof ConstTest' must have a '[Symbol.iterator]()' method that returns an iterator.


==== tests/cases/compiler/enumIteratorConstEnum1.ts (2 errors) ====
const enum ConstTest {
Foo,
Bar,
Baz,
}

// Should be an error.
for (const member of ConstTest) {
~~~~~~~~~
!!! error TS2475: 'const' enums can only be used in property or index access expressions or the right hand side of an import declaration or export assignment or type query.
~~~~~~~~~
!!! error TS2488: Type 'typeof ConstTest' must have a '[Symbol.iterator]()' method that returns an iterator.
}

16 changes: 16 additions & 0 deletions tests/baselines/reference/enumIteratorConstEnum1.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
//// [enumIteratorConstEnum1.ts]
const enum ConstTest {
Foo,
Bar,
Baz,
}

// Should be an error.
for (const member of ConstTest) {
}


//// [enumIteratorConstEnum1.js]
// Should be an error.
for (const member of ConstTest) {
}
20 changes: 20 additions & 0 deletions tests/baselines/reference/enumIteratorConstEnum1.symbols
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
=== tests/cases/compiler/enumIteratorConstEnum1.ts ===
const enum ConstTest {
>ConstTest : Symbol(ConstTest, Decl(enumIteratorConstEnum1.ts, 0, 0))

Foo,
>Foo : Symbol(ConstTest.Foo, Decl(enumIteratorConstEnum1.ts, 0, 22))

Bar,
>Bar : Symbol(ConstTest.Bar, Decl(enumIteratorConstEnum1.ts, 1, 6))

Baz,
>Baz : Symbol(ConstTest.Baz, Decl(enumIteratorConstEnum1.ts, 2, 6))
}

// Should be an error.
for (const member of ConstTest) {
>member : Symbol(member, Decl(enumIteratorConstEnum1.ts, 7, 10))
>ConstTest : Symbol(ConstTest, Decl(enumIteratorConstEnum1.ts, 0, 0))
}

20 changes: 20 additions & 0 deletions tests/baselines/reference/enumIteratorConstEnum1.types
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
=== tests/cases/compiler/enumIteratorConstEnum1.ts ===
const enum ConstTest {
>ConstTest : ConstTest

Foo,
>Foo : ConstTest.Foo

Bar,
>Bar : ConstTest.Bar

Baz,
>Baz : ConstTest.Baz
}

// Should be an error.
for (const member of ConstTest) {
>member : any
>ConstTest : typeof ConstTest
}

17 changes: 17 additions & 0 deletions tests/cases/compiler/enumIteratorBasics1.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
// @target: ES6
enum Test {
Foo,
Bar,
Baz,
}

const AlsoTest = Test;

for (const member of AlsoTest) {
const x: string = member[0];
const y: Test = member[1];
console.log(x, y);
}

const members: [string, Test][] = [...AlsoTest];
console.log(members);
Loading