Support enums.

Author: milesj

src/convertBabelToPropTypes.ts

@@ -8,6 +8,7 @@
 import { types as t } from '@babel/core';
 import { addComment } from '@babel/types';
 import { convertSymbolFromSource } from './convertTSToPropTypes';
+import extractEnumValues from './extractEnumValues';
 import getTypeName from './getTypeName';
 import {
   createCall,
@@ -75,6 +76,24 @@ function convert(type: any, state: ConvertState, depth: number): PropType | null
       propTypesImportedName,
     );
 
+    // enum Foo {} -> PropTypes.oneOf
+  } else if (t.isTSEnumDeclaration(type)) {
+    return createCall(
+      t.identifier('oneOf'),
+      [
+        t.arrayExpression(
+          extractEnumValues(type).map(value => {
+            if (typeof value === 'number') {
+              return t.numericLiteral(value);
+            }
+
+            return t.stringLiteral(value);
+          }),
+        ),
+      ],
+      propTypesImportedName,
+    );
+
     // (() => void) -> PropTypes.func
   } else if (t.isTSFunctionType(type)) {
     return createMember(t.identifier('func'), propTypesImportedName);

src/extractEnumValues.ts

@@ -0,0 +1,24 @@
+import { types as t } from '@babel/core';
+
+export default function extractEnumValues(enumDecl: t.TSEnumDeclaration): (string | number)[] {
+  const values: (string | number)[] = [];
+  let lastIndex = -1;
+
+  enumDecl.members.forEach(({ initializer }) => {
+    if (initializer) {
+      if (t.isNumericLiteral(initializer)) {
+        lastIndex = initializer.value;
+        values.push(initializer.value);
+      } else if (t.isStringLiteral(initializer)) {
+        values.push(initializer.value);
+      } else {
+        // Skip
+      }
+    } else {
+      lastIndex += 1;
+      values.push(lastIndex);
+    }
+  });
+
+  return values;
+}

src/index.ts

@@ -232,6 +232,11 @@ export default declare((api: any, options: PluginOptions, root: string) => {
               }
             },
 
+            // `enum Foo {}`
+            TSEnumDeclaration({ node }: Path<t.TSInterfaceDeclaration>) {
+              state.referenceTypes[node.id.name] = node;
+            },
+
             // `interface FooProps {}`
             TSInterfaceDeclaration({ node }: Path<t.TSInterfaceDeclaration>) {
               state.componentTypes[node.id.name] = extractTypeProperties(

src/types.ts

@@ -40,7 +40,9 @@ export interface ConvertState {
     hasImport: boolean;
   };
   reactImportedName: string;
-  referenceTypes: { [key: string]: t.TSInterfaceDeclaration | t.TSTypeAliasDeclaration };
+  referenceTypes: {
+    [key: string]: t.TSEnumDeclaration | t.TSInterfaceDeclaration | t.TSTypeAliasDeclaration;
+  };
   typeChecker?: ts.TypeChecker;
   typeProgram?: ts.Program;
 }

tests/__snapshots__/index.test.ts.snap

@@ -1216,6 +1216,65 @@ export default class TypeComplex extends React.Component<Props> {
 }"
 `;
 
+exports[`babel-plugin-typescript-to-proptypes transforms ./fixtures/type-enum.ts 1`] = `
+"import _pt from 'prop-types';
+import React from 'react';
+enum Standard {
+  FOO,
+  BAR,
+  BAZ,
+}
+enum FirstInitializer {
+  FOO = 10,
+  BAR,
+  BAZ,
+}
+enum AllNumbers {
+  FOO = 10,
+  BAR = 20,
+  BAZ = 30,
+}
+enum AllStrings {
+  FOO = 'Foo',
+  BAR = 'Bar',
+  BAZ = 'Baz',
+}
+enum MixedTypes {
+  FOO = 1,
+  BAR = 'Bar',
+  BAZ = 2,
+}
+const enum ConstEnum {
+  Up,
+  Down,
+  Left,
+  Right,
+}
+export interface Props {
+  standard: Standard;
+  first?: FirstInitializer;
+  numbers?: AllNumbers;
+  strings: AllStrings;
+  mixed?: MixedTypes;
+  constant?: ConstEnum;
+}
+export default class TypeEnum extends React.Component<Props> {
+  static propTypes = {
+    standard: _pt.oneOf([0, 1, 2]).isRequired,
+    first: _pt.oneOf([10, 11, 12]),
+    numbers: _pt.oneOf([10, 20, 30]),
+    strings: _pt.oneOf(['Foo', 'Bar', 'Baz']).isRequired,
+    mixed: _pt.oneOf([1, 'Bar', 2]),
+    constant: _pt.oneOf([0, 1, 2, 3])
+  };
+
+  render() {
+    return null;
+  }
+
+}"
+`;
+
 exports[`babel-plugin-typescript-to-proptypes transforms ./fixtures/type-function.ts 1`] = `
 "import _pt from 'prop-types';
 import React from 'react';

tests/fixtures/type-enum.ts

@@ -0,0 +1,53 @@
+import React from 'react';
+
+enum Standard {
+  FOO,
+  BAR,
+  BAZ,
+}
+
+enum FirstInitializer {
+  FOO = 10,
+  BAR,
+  BAZ,
+}
+
+enum AllNumbers {
+  FOO = 10,
+  BAR = 20,
+  BAZ = 30,
+}
+
+enum AllStrings {
+  FOO = 'Foo',
+  BAR = 'Bar',
+  BAZ = 'Baz',
+}
+
+enum MixedTypes {
+  FOO = 1,
+  BAR = 'Bar',
+  BAZ = 2,
+}
+
+const enum ConstEnum {
+  Up,
+  Down,
+  Left,
+  Right,
+}
+
+export interface Props {
+  standard: Standard;
+  first?: FirstInitializer;
+  numbers?: AllNumbers;
+  strings: AllStrings;
+  mixed?: MixedTypes;
+  constant?: ConstEnum;
+}
+
+export default class TypeEnum extends React.Component<Props> {
+  render() {
+    return null;
+  }
+}