Add support for @enum to for all variables with string literal prop…

…erties Resolves #1740.
TypeStrong · Oct 17, 2021 · 11e7b4f · 11e7b4f
1 parent c2f9a89
commit 11e7b4f
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ### Features
 
 -   Added support for displaying identifiers & property access expressions in initializers, #1730.
+-   Expanded support for variables tagged with `@enum` to all variables whose property types are string literals, #1740.
 
 ### Bug Fixes
 

diff --git a/src/lib/converter/symbols.ts b/src/lib/converter/symbols.ts
@@ -818,7 +818,7 @@ function convertVariable(
     const type = context.checker.getTypeOfSymbolAtLocation(symbol, declaration);
 
     if (
-        isEnumLiteral(declaration as ts.VariableDeclaration) &&
+        isEnumLike(context.checker, type, declaration) &&
         symbol.getJsDocTags().some((tag) => tag.name === "enum")
     ) {
         return convertVariableAsEnum(context, symbol, exportSymbol);
@@ -852,27 +852,15 @@ function convertVariable(
     context.finalizeDeclarationReflection(reflection, symbol, exportSymbol);
 }
 
-function isEnumLiteral(declaration: ts.VariableDeclaration) {
-    if (
-        !hasAllFlags(declaration.parent.flags, ts.NodeFlags.Const) ||
-        !declaration.initializer ||
-        !ts.isAsExpression(declaration.initializer) ||
-        declaration.initializer.type.getText() !== "const"
-    ) {
-        // Not an as-const expression
-        return false;
-    }
-
-    const body = declaration.initializer.expression;
-    if (!ts.isObjectLiteralExpression(body)) {
+function isEnumLike(checker: ts.TypeChecker, type: ts.Type, location: ts.Node) {
+    if (!hasAllFlags(type.flags, ts.TypeFlags.Object)) {
         return false;
     }
 
-    return body.properties.every(
-        (prop) =>
-            ts.isPropertyAssignment(prop) &&
-            ts.isStringLiteral(prop.initializer)
-    );
+    return type.getProperties().every((prop) => {
+        const propType = checker.getTypeOfSymbolAtLocation(prop, location);
+        return propType.isStringLiteral();
+    });
 }
 
 function convertVariableAsEnum(
@@ -889,22 +877,24 @@ function convertVariableAsEnum(
     const rc = context.withScope(reflection);
 
     const declaration = symbol.declarations![0] as ts.VariableDeclaration;
-    const init = (declaration.initializer as ts.AsExpression)
-        .expression as ts.ObjectLiteralExpression;
+    const type = context.checker.getTypeAtLocation(declaration);
 
-    for (const prop of init.properties as readonly ts.PropertyAssignment[]) {
-        const childSymbol = context.checker.getSymbolAtLocation(prop.name);
+    for (const prop of type.getProperties()) {
         const reflection = rc.createDeclarationReflection(
             ReflectionKind.EnumMember,
-            childSymbol,
+            prop,
             void 0
         );
 
-        reflection.defaultValue = JSON.stringify(
-            (prop.initializer as ts.StringLiteral).text
+        const propType = context.checker.getTypeOfSymbolAtLocation(
+            prop,
+            declaration
         );
+        assert(propType.isStringLiteral());
+
+        reflection.defaultValue = JSON.stringify(propType.value);
 
-        rc.finalizeDeclarationReflection(reflection, childSymbol, void 0);
+        rc.finalizeDeclarationReflection(reflection, prop, void 0);
     }
 
     // Skip converting the type alias, if there is one

diff --git a/src/test/behaviorTests.ts b/src/test/behaviorTests.ts
@@ -17,11 +17,28 @@ export const behaviorTests: Record<
 > = {
     asConstEnum(project) {
         const SomeEnumLike = query(project, "SomeEnumLike");
-        equal(SomeEnumLike.kind, ReflectionKind.Variable);
+        equal(SomeEnumLike.kind, ReflectionKind.Variable, "SomeEnumLike");
         const SomeEnumLikeTagged = query(project, "SomeEnumLikeTagged");
-        equal(SomeEnumLikeTagged.kind, ReflectionKind.Enum);
+        equal(
+            SomeEnumLikeTagged.kind,
+            ReflectionKind.Enum,
+            "SomeEnumLikeTagged"
+        );
         const A = query(project, "SomeEnumLikeTagged.a");
         equal(A.defaultValue, '"a"');
+
+        const ManualEnum = query(project, "ManualEnum");
+        equal(ManualEnum.kind, ReflectionKind.Enum, "ManualEnum");
+
+        const ManualWithoutHelper = query(project, "ManualEnumHelper");
+        equal(
+            ManualWithoutHelper.kind,
+            ReflectionKind.Enum,
+            "ManualEnumHelper"
+        );
+
+        const WithoutReadonly = query(project, "WithoutReadonly");
+        equal(WithoutReadonly.kind, ReflectionKind.Enum, "WithoutReadonly");
     },
     duplicateHeritageClauses(project) {
         const b = query(project, "B");

diff --git a/src/test/converter2/behavior/asConstEnum.ts b/src/test/converter2/behavior/asConstEnum.ts
@@ -8,3 +8,18 @@ export const SomeEnumLikeTagged = {
     a: "a",
     b: "b",
 } as const;
+
+/** @enum */
+export const ManualEnum: { readonly a: "a" } = {
+    a: "a",
+};
+
+/** @enum */
+export const ManualEnumHelper: Readonly<{ a: "a" }> = {
+    a: "a",
+};
+
+/** @enum */
+export const WithoutReadonly = {
+    a: "a",
+} as { a: "a" };