Skip to content

Commit e801cd9

Browse files
author
Tim Hutt
committed
Add support for iterating over enums
I'm fairly confident about the codegen stuff, but I have no idea about the type checker. I also couldn't make it work with type aliases - it complains that the alias is only a type but you're trying to use it as a value.
1 parent 4c62f6e commit e801cd9

File tree

2 files changed

+58
-0
lines changed

2 files changed

+58
-0
lines changed

src/compiler/checker.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35044,6 +35044,17 @@ namespace ts {
3504435044
return anyIterationTypes;
3504535045
}
3504635046

35047+
const symbolType = getDeclaredTypeOfSymbol(type.symbol);
35048+
35049+
if (symbolType.symbol.flags & SymbolFlags.RegularEnum) {
35050+
return createIterationTypes(
35051+
createTupleType([
35052+
stringType,
35053+
symbolType,
35054+
]),
35055+
);
35056+
}
35057+
3504735058
if (use & IterationUse.AllowsAsyncIterablesFlag) {
3504835059
const iterationTypes =
3504935060
getIterationTypesOfIterableCached(type, asyncIterationTypesResolver) ||

src/compiler/transformers/ts.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2426,6 +2426,7 @@ namespace ts {
24262426
const statements: Statement[] = [];
24272427
startLexicalEnvironment();
24282428
const members = map(node.members, transformEnumMember);
2429+
members.push(transformEnumMembersToIterator(node.members));
24292430
insertStatementsAfterStandardPrologue(statements, endLexicalEnvironment());
24302431
addRange(statements, members);
24312432

@@ -2474,6 +2475,52 @@ namespace ts {
24742475
);
24752476
}
24762477

2478+
/**
2479+
* Transforms all enum members into a generator function that yields
2480+
* the enum members.
2481+
*
2482+
* @param members The enum member nodes.
2483+
*/
2484+
function transformEnumMembersToIterator(members: readonly EnumMember[]): Statement {
2485+
const symbolIterator = factory.createPropertyAccessExpression(
2486+
factory.createIdentifier("Symbol"),
2487+
"iterator",
2488+
);
2489+
const iteratorMember = factory.createElementAccessExpression(
2490+
currentNamespaceContainerName,
2491+
symbolIterator,
2492+
);
2493+
2494+
const generatorBody = factory.createBlock(members.map(member => {
2495+
// Todo: don't call these twice.
2496+
const name = getExpressionForPropertyName(member, /*generateNameForComputedPropertyName*/ false);
2497+
const valueExpression = transformEnumMemberDeclarationValue(member);
2498+
const yieldExpression = factory.createYieldExpression(
2499+
undefined,
2500+
factory.createArrayLiteralExpression([
2501+
name,
2502+
valueExpression,
2503+
]),
2504+
);
2505+
return factory.createExpressionStatement(yieldExpression);
2506+
}));
2507+
2508+
const iteratorGenerator = factory.createFunctionExpression(
2509+
undefined,
2510+
factory.createToken(SyntaxKind.AsteriskToken),
2511+
undefined,
2512+
undefined,
2513+
[],
2514+
undefined,
2515+
generatorBody,
2516+
);
2517+
const outerAssignment = factory.createAssignment(
2518+
iteratorMember,
2519+
iteratorGenerator,
2520+
);
2521+
return factory.createExpressionStatement(outerAssignment);
2522+
}
2523+
24772524
/**
24782525
* Transforms the value of an enum member.
24792526
*

0 commit comments

Comments
 (0)