Merge pull request #3 from IntrinsicLabsAI/aduffy/tweak-enum

Small tweaks on top of #2

.github/workflows/npm.yml

@@ -13,7 +13,7 @@ jobs:
           node-version: '16.x'
           registry-url: 'https://registry.npmjs.org'
       - run: npm ci
-      - run: npx tsc
+      - run: npm run build
       - run: npm --no-git-tag-version version from-git
       - run: npm publish --access public
         env:

README.md

@@ -25,12 +25,20 @@ import { compile, serializeGrammar } from "@intrinsicai/gbnfgen";
 
 // Supporting Enum for multiple choices (cannot be numbers)
 const grammar = compile(
-    `enum Mood { Happy, Sad, Grateful, Excited, Angry, Peaceful }
+    `enum Mood {
+      Happy = "happy",
+      Sad = "sad",
+      Grateful = "grateful",
+      Excited = "excited", 
+      Angry = "angry, 
+      Peaceful = "peaceful"
+    }
     
     interface Person {
          name: string;
          occupation: string;
          age: number;
+         mood: Mood,
      }`, "Person");
 ```
 

src/compiler.test.ts

@@ -30,7 +30,7 @@ numberlist ::= "["   ws   "]" | "["   ws   string   (","   ws   number)*   ws
 
 test("Single interface with enum generation", () => {
   const postalAddressGrammar = compile(
-    `enum AddressType { business, home };
+    `enum AddressType { Business = "business", Home = "home" };
     interface PostalAddress {
     streetNumber: number;
     type: AddressType;
@@ -41,21 +41,21 @@ test("Single interface with enum generation", () => {
   }`,
     "PostalAddress"
   );
-
-
+
   expect(serializeGrammar(postalAddressGrammar).trimEnd()).toEqual(
     String.raw`
 root ::= PostalAddress
-PostalAddress ::= "{"   ws   "\"streetNumber\":"   ws   number   ","   ws   "\"type\":"   ws   enumAddressType   ","   ws   "\"street\":"   ws   string   ","   ws   "\"city\":"   ws   string   ","   ws   "\"state\":"   ws   string   ","   ws   "\"postalCode\":"   ws   number   "}"
+PostalAddress ::= "{"   ws   "\"streetNumber\":"   ws   number   ","   ws   "\"type\":"   ws   AddressType   ","   ws   "\"street\":"   ws   string   ","   ws   "\"city\":"   ws   string   ","   ws   "\"state\":"   ws   string   ","   ws   "\"postalCode\":"   ws   number   "}"
 PostalAddresslist ::= "[]" | "["   ws   PostalAddress   (","   ws   PostalAddress)*   "]"
+AddressType ::= "\"" "business" "\"" | "\"" "home" "\""
 string ::= "\""   ([^"]*)   "\""
 boolean ::= "true" | "false"
 ws ::= [ \t\n]*
 number ::= [0-9]+   "."?   [0-9]*
 stringlist ::= "["   ws   "]" | "["   ws   string   (","   ws   string)*   ws   "]"
 numberlist ::= "["   ws   "]" | "["   ws   string   (","   ws   number)*   ws   "]"
-enumAddressType ::= "\"" "business" "\"" | "\"" "home" "\""`.trim()
-  )
+`.trim()
+  );
 });
 
 test("Single multiple interface with references generation", () => {
@@ -96,9 +96,9 @@ test("Single multiple interface and enum with references generation", () => {
     `
     // Define an enum for product categories
     enum ProductCategory {
-      Electronics,
-      Clothing,
-      Food
+      Electronics = "Electronics",
+      Clothing = "Clothing",
+      Food = "Food"
     }
     
     // Define an interface for representing a product
@@ -112,10 +112,10 @@ test("Single multiple interface and enum with references generation", () => {
     
     // Define an enum for order statuses
     enum OrderStatus {
-      Pending,
-      Shipped,
-      Delivered,
-      Canceled
+      Pending = "Pending",
+      Shipped = "Shipped",
+      Delivered = "Delivered",
+      Canceled = "Canceled"
     }
     
     // Define an interface for representing an order
@@ -128,24 +128,24 @@ test("Single multiple interface and enum with references generation", () => {
   `,
     "Order"
   );
-
 
   expect(serializeGrammar(resumeGrammar).trimEnd()).toEqual(
     String.raw`
 root ::= Order
-Order ::= "{"   ws   "\"orderId\":"   ws   number   ","   ws   "\"products\":"   ws   Productlist   ","   ws   "\"status\":"   ws   enumOrderStatus   ","   ws   "\"orderDate\":"   ws   string   "}"
+Order ::= "{"   ws   "\"orderId\":"   ws   number   ","   ws   "\"products\":"   ws   Productlist   ","   ws   "\"status\":"   ws   OrderStatus   ","   ws   "\"orderDate\":"   ws   string   "}"
 Orderlist ::= "[]" | "["   ws   Order   (","   ws   Order)*   "]"
-Product ::= "{"   ws   "\"id\":"   ws   number   ","   ws   "\"name\":"   ws   string   ","   ws   "\"description\":"   ws   string   ","   ws   "\"price\":"   ws   number   ","   ws   "\"category\":"   ws   enumProductCategory   "}"
+OrderStatus ::= "\"" "Pending" "\"" | "\"" "Shipped" "\"" | "\"" "Delivered" "\"" | "\"" "Canceled" "\""
+Product ::= "{"   ws   "\"id\":"   ws   number   ","   ws   "\"name\":"   ws   string   ","   ws   "\"description\":"   ws   string   ","   ws   "\"price\":"   ws   number   ","   ws   "\"category\":"   ws   ProductCategory   "}"
 Productlist ::= "[]" | "["   ws   Product   (","   ws   Product)*   "]"
+ProductCategory ::= "\"" "Electronics" "\"" | "\"" "Clothing" "\"" | "\"" "Food" "\""
 string ::= "\""   ([^"]*)   "\""
 boolean ::= "true" | "false"
 ws ::= [ \t\n]*
 number ::= [0-9]+   "."?   [0-9]*
 stringlist ::= "["   ws   "]" | "["   ws   string   (","   ws   string)*   ws   "]"
 numberlist ::= "["   ws   "]" | "["   ws   string   (","   ws   number)*   ws   "]"
-enumProductCategory ::= "\"" "Electronics" "\"" | "\"" "Clothing" "\"" | "\"" "Food" "\""
-enumOrderStatus ::= "\"" "Pending" "\"" | "\"" "Shipped" "\"" | "\"" "Delivered" "\"" | "\"" "Canceled" "\""`.trim()
-  )
+`.trim()
+  );
 });
 
 test("Jsonformer car example", () => {

src/compiler.ts

@@ -10,7 +10,14 @@ import {
   sequence,
 } from "./grammar.js";
 
-import { toElementId, toListElementId, WS_REF, getGrammarRegister, GrammarRegister, registerToGrammar } from "./util.js";
+import {
+  toElementId,
+  toListElementId,
+  WS_REF,
+  getDefaultGrammar,
+  GrammarRegister,
+  registerToGrammar,
+} from "./util.js";
 
 // Turn interface properties into Grammar References
 export function toGrammar(iface: Interface): Grammar {
@@ -68,8 +75,7 @@ export function toGrammar(iface: Interface): Grammar {
 }
 
 // Parameterized list of things
-export type PropertyType = string
-  | { reference: string; isArray: boolean };
+export type PropertyType = string | { reference: string; isArray: boolean };
 
 export interface InterfaceProperty {
   name: string;
@@ -123,13 +129,20 @@ function handleEnum(enumNode: EnumDeclaration): GrammarElement {
   const choices: GrammarRule[] = [];
   if (enumNode && enumNode.members) {
     for (const member of enumNode.members) {
-      if (ts.isEnumMember(member) && member.name && ts.isIdentifier(member.name)) {
-        choices.push(literal(member.name.text, true));
+      // NOTE(aduffy): support union type literals as well.
+      if (ts.isEnumMember(member) && ts.isIdentifier(member.name)) {
+        // If initializer is String, we use the string value. Else, we assume a numeric value.
+        if (!member.initializer || !ts.isStringLiteral(member.initializer)) {
+          throw new Error(
+            "Only string enums are supported. Please check the String enums section of the TypeScript Handbook at https://www.typescriptlang.org/docs/handbook/enums.html"
+          );
+        }
+        choices.push(literal(member.initializer.text, true));
       }
     }
   }
 
-  return { identifier: `enum${enumNode.name.text}`, alternatives: choices };
+  return { identifier: enumNode.name.text, alternatives: choices };
 }
 
 function handleInterface(
@@ -159,13 +172,11 @@ function handleInterface(
     }
     const propName = child.name.getText(srcFile);
     const propType = child.type?.getText(srcFile) ?? "never";
-  
+
     // Validate one of the accepted types
     let propTypeValidated: PropertyType;
     if (register.has(propType)) {
       propTypeValidated = propType;
-    } else if (register.has(`enum${propType}`)) {
-      propTypeValidated = `enum${propType}`;
     } else if (propType === "string[]" || propType === "Array<string>") {
       propTypeValidated = "stringlist";
     } else if (propType === "number[]" || propType === "Array<number>") {
@@ -212,7 +223,7 @@ export function compile(source: string, rootType: string): Grammar {
   });
 
   // Get the default Grammar Register
-  const register = getGrammarRegister();
+  const register = getDefaultGrammar();
 
   // Run the compiler to ensure that the typescript source file is correct.
   const emitResult = program.emit();
@@ -233,10 +244,9 @@ export function compile(source: string, rootType: string): Grammar {
       declaredTypes.add(child.name.getText(srcFile));
     }
 
-    // Add the Enum to Gramma Register
+    // Add the Enum to Grammar Register
     if (ts.isEnumDeclaration(child)) {
-      const element = handleEnum(child);
-      register.set(element.identifier, element.alternatives);
+      declaredTypes.add(child.name.getText(srcFile));
     }
   });
 
@@ -247,25 +257,19 @@ export function compile(source: string, rootType: string): Grammar {
     );
   }
 
-  // Create the Enum Type for each enum
-
-  // Define basic grammar rules
+  // Import default grammar rules
   const grammar: Grammar = {
-    elements: [...registerToGrammar(register)]
-  }
+    elements: [...registerToGrammar(register)],
+  };
 
   srcFile.forEachChild((child) => {
     if (ts.isInterfaceDeclaration(child)) {
-      const iface = handleInterface(
-        child,
-        srcFile,
-        declaredTypes,
-        register
-      );
+      const iface = handleInterface(child, srcFile, declaredTypes, register);
       const ifaceGrammar = toGrammar(iface);
-
-      // Add grammar rules above basic grammar rules
       grammar.elements.unshift(...ifaceGrammar.elements);
+    } else if (ts.isEnumDeclaration(child)) {
+      const enumGrammar = handleEnum(child);
+      grammar.elements.unshift(enumGrammar);
     }
   });
 

src/grammar.ts

@@ -82,16 +82,18 @@ function serializeSequence(rule: RuleSequence): string {
 
 function serializeGroup(rule: RuleGroup): string {
   const multiplicity = {
-    "none": "",
-    "optional": "?",
-    "star": "*",
-    "plus": "+",
+    none: "",
+    optional: "?",
+    star: "*",
+    plus: "+",
   }[rule.multiplicity];
   return `(${serializeSequence(rule.rules)})${multiplicity}`;
 }
 
 function serializeLiteralRule(rule: RuleLiteral): string {
-  return rule.quote ? "\"\\\"\" " + JSON.stringify(rule.literal) + " \"\\\"\"" : JSON.stringify(rule.literal);
+  return rule.quote
+    ? '"\\"" ' + JSON.stringify(rule.literal) + ' "\\""'
+    : JSON.stringify(rule.literal);
 }
 
 function serializeReference(rule: RuleReference): string {
@@ -136,7 +138,7 @@ export function serializeGrammar(grammar: Grammar): string {
   return out;
 }
 
-export function literal(value: string, quote: boolean=false): RuleLiteral {
+export function literal(value: string, quote: boolean = false): RuleLiteral {
   return {
     type: "literal",
     literal: value,
@@ -165,10 +167,13 @@ export function reference(value: string): RuleReference {
   };
 }
 
-export function group(rules: RuleSequence, multiplicity: RuleGroup["multiplicity"]): RuleGroup {
+export function group(
+  rules: RuleSequence,
+  multiplicity: RuleGroup["multiplicity"]
+): RuleGroup {
   return {
     type: "group",
     rules,
     multiplicity,
-  }
-}
+  };
+}

src/util.ts

@@ -85,7 +85,7 @@ export type GrammarRegister = Map<string, Array<GrammarRule>>;
  * Enables Enum add to the register when compiling the source file.
  * @returns The Default Grammar Element Register
  */
-export function getGrammarRegister(): GrammarRegister {
+export function getDefaultGrammar(): GrammarRegister {
   const register = new Map<string, Array<GrammarRule>>();
 
   register.set(STRING_ELEM.identifier, STRING_ELEM.alternatives);