Merge pull request #3552 from trufflesuite/usenums

Enhancement: Add enum objects to contract objects

packages/contract-tests/test/enums.js

@@ -0,0 +1,22 @@
+var assert = require("chai").assert;
+var util = require("./util");
+
+describe("Enumerations", function () {
+  var Example;
+
+  before(async function () {
+    this.timeout(10000);
+
+    Example = await util.createExample();
+  });
+
+  it("Sets up enumeration objects on contracts", async function () {
+    const expected = {
+      ExampleZero: 0,
+      ExampleOne: 1,
+      ExampleTwo: 2
+    };
+    assert.deepEqual(Example.enums.ExampleEnum, expected);
+    assert.deepEqual(Example.ExampleEnum, expected);
+  });
+});

packages/contract-tests/test/sources/Example.sol

@@ -13,6 +13,8 @@ contract Example {
   event SpecialEvent();
   event NumberEvent(int numA, int indexed numB, address addrC, uint numD, uint);
 
+  enum ExampleEnum { ExampleZero, ExampleOne, ExampleTwo }
+
   constructor(uint val) public {
     // Constructor revert
     require(val != 13);

packages/contract/lib/contract/bootstrap.js

@@ -1,4 +1,5 @@
 const execute = require("../execute");
+const debug = require("debug")("contract:contract:bootstrap");
 
 module.exports = fn => {
   // Add our static methods
@@ -16,5 +17,50 @@ module.exports = fn => {
   fn["new"].estimateGas = execute.estimateDeployment.bind(fn);
   fn["new"].request = execute.requestDeployment.bind(fn);
 
+  //add enumerations. (probably these should go in
+  //constructorMethods.js, but this is easier to modify... we'll
+  //redo all this in the rewrite anyway)
+  if (fn._json) {
+    //getters will throw otherwise!
+    if (fn.ast) {
+      //note this was set up earlier
+      const node = locateNode(fn.contractName, fn.ast); //name also set up earlier
+      if (node) {
+        fn.enums = extractEnums(node);
+        for (const [name, enumeration] of Object.entries(fn.enums)) {
+          //enum is a reserved word :P
+          if (!(name in fn)) {
+            //don't overwrite anything!
+            fn[name] = enumeration;
+          }
+        }
+      }
+    }
+  }
+
   return fn;
 };
+
+function locateNode(name, ast) {
+  if (ast.nodeType === "SourceUnit") {
+    return ast.nodes.find(
+      node => node.nodeType === "ContractDefinition" && node.name === name
+    );
+  } else {
+    return undefined;
+  }
+}
+
+function extractEnums(node) {
+  return Object.assign(
+    {},
+    ...node.nodes
+      .filter(definition => definition.nodeType === "EnumDefinition")
+      .map(definition => ({
+        [definition.name]: Object.assign(
+          {},
+          ...definition.members.map((member, index) => ({[member.name]: index}))
+        )
+      }))
+  );
+}

packages/contract/lib/contract/index.js

@@ -1,4 +1,4 @@
-const debug = require("debug")("contract:contract"); // eslint-disable-line no-unused-vars
+const debug = require("debug")("contract:contract");
 let Web3 = require("web3");
 const webUtils = require("web3-utils");
 const execute = require("../execute");
@@ -12,7 +12,7 @@ if (typeof Web3 === "object" && Object.keys(Web3).length === 0) {
   Web3 = global.Web3;
 }
 
-(function(module) {
+(function (module) {
   // Accepts a contract object created with web3.eth.Contract or an address.
   function Contract(contract) {
     var instance = this;
@@ -38,15 +38,15 @@ if (typeof Web3 === "object" && Object.keys(Web3).length === 0) {
     }
 
     // User defined methods, overloaded methods, events
-    instance.abi.forEach(function(item) {
+    instance.abi.forEach(function (item) {
       switch (item.type) {
         case "function":
           var isConstant =
             ["pure", "view"].includes(item.stateMutability) || item.constant; // new form // deprecated case
 
           var signature = webUtils._jsonInterfaceMethodToString(item);
 
-          var method = function(constant, web3Method) {
+          var method = function (constant, web3Method) {
             var fn;
 
             constant
@@ -127,7 +127,7 @@ if (typeof Web3 === "object" && Object.keys(Web3).length === 0) {
     // Prefer user defined `send`
     if (!instance.send) {
       instance.send = (value, txParams = {}) => {
-        const packet = Object.assign({ value: value }, txParams);
+        const packet = Object.assign({value: value}, txParams);
         return instance.sendTransaction(packet);
       };
     }