feat: allow named types in unions

lib/types.js

@@ -111,6 +111,8 @@ class Type {
         wrapUnions = 'auto';
       } else if (typeof wrapUnions == 'string') {
         wrapUnions = wrapUnions.toLowerCase();
+      } else if (typeof wrapUnions === 'function') {
+        wrapUnions = 'auto';
       }
       switch (wrapUnions) {
         case 'always':
@@ -196,11 +198,26 @@ class Type {
       let types = schema.map((obj) => {
         return Type.forSchema(obj, opts);
       });
+      let projectionFn;
       if (!UnionType) {
-        UnionType = isAmbiguous(types) ? WrappedUnionType : UnwrappedUnionType;
+        // either automatic detection or we have a projection function
+        if (typeof opts.wrapUnions === 'function') {
+          // we have a projection function
+          try {
+            projectionFn = opts.wrapUnions(types);
+            UnionType = typeof projectionFn !== 'undefined'
+              // projection function yields a function, we can use an Unwrapped type
+              ? UnwrappedUnionType
+              : WrappedUnionType;
+          } catch(e) {
+            throw new Error(`Error generating projection function: ${e}`);
+          }
+        } else {
+          UnionType = isAmbiguous(types) ? WrappedUnionType : UnwrappedUnionType;
+        }
       }
       LOGICAL_TYPE = logicalType;
-      type = new UnionType(types, opts);
+      type = new UnionType(types, opts, projectionFn);
     } else { // New type definition.
       type = (function (typeName) {
         let Type = TYPES[typeName];
@@ -341,10 +358,10 @@ class Type {
           return branchTypes[name];
         }), opts);
       } catch (err) {
-        opts.wrapUnions = wrapUnions;
         throw err;
+      } finally {
+        opts.wrapUnions = wrapUnions;
       }
-      opts.wrapUnions = wrapUnions;
       return unionType;
     }
 
@@ -1226,6 +1243,44 @@ UnionType.prototype._branchConstructor = function () {
   throw new Error('unions cannot be directly wrapped');
 };
 
+
+function generateProjectionIndexer(projectionFn) {
+  return (val) => {
+    const index = projectionFn(val);
+    if (typeof index !== 'number') {
+      throw new Error(`Projected index '${index}' is not valid`);
+    }
+    return index;
+  };
+}
+
+function generateDefaultIndexer() {
+  this._dynamicBranches = null;
+  this._bucketIndices = {};
+  this.types.forEach(function (type, index) {
+    if (Type.isType(type, 'abstract', 'logical')) {
+      if (!this._dynamicBranches) {
+        this._dynamicBranches = [];
+      }
+      this._dynamicBranches.push({index, type});
+    } else {
+      let bucket = getTypeBucket(type);
+      if (this._bucketIndices[bucket] !== undefined) {
+        throw new Error(`ambiguous unwrapped union: ${j(this)}`);
+      }
+      this._bucketIndices[bucket] = index;
+    }
+  }, this);
+  return (val) => {
+    let index = this._bucketIndices[getValueBucket(val)];
+    if (this._dynamicBranches) {
+      // Slower path, we must run the value through all branches.
+      index = this._getBranchIndex(val, index);
+    }
+    return index;
+  };
+}
+
 /**
  * "Natural" union type.
  *
@@ -1246,38 +1301,33 @@ UnionType.prototype._branchConstructor = function () {
  * + `map`, `record`
  */
 class UnwrappedUnionType extends UnionType {
-  constructor (schema, opts) {
+  /**
+   *
+   * @param {*} schema
+   * @param {*} opts
+   * @param {Function|undefined} projectionFn The projection function used
+   *                                          to determine the bucket for the
+   *                                          Union. Falls back to generate
+   *                                          from `wrapUnions` parameter
+   *                                          if given.
+   */
+  constructor (schema, opts, projectionFn) {
     super(schema, opts);
 
-    this._dynamicBranches = null;
-    this._bucketIndices = {};
-    this.types.forEach(function (type, index) {
-      if (Type.isType(type, 'abstract', 'logical')) {
-        if (!this._dynamicBranches) {
-          this._dynamicBranches = [];
-        }
-        this._dynamicBranches.push({index, type});
-      } else {
-        let bucket = getTypeBucket(type);
-        if (this._bucketIndices[bucket] !== undefined) {
-          throw new Error(`ambiguous unwrapped union: ${j(this)}`);
-        }
-        this._bucketIndices[bucket] = index;
+    if (!projectionFn && opts && typeof opts.wrapUnions === 'function') {
+      try {
+        projectionFn = opts.wrapUnions(this.types);
+      } catch(e) {
+        throw new Error(`Error generating projection function: ${e}`);
       }
-    }, this);
+    }
+    this._getIndex = projectionFn
+      ? generateProjectionIndexer(projectionFn)
+      : generateDefaultIndexer.bind(this)(this.types);
 
     Object.freeze(this);
   }
 
-  _getIndex (val) {
-    let index = this._bucketIndices[getValueBucket(val)];
-    if (this._dynamicBranches) {
-      // Slower path, we must run the value through all branches.
-      index = this._getBranchIndex(val, index);
-    }
-    return index;
-  }
-
   _getBranchIndex (any, index) {
     let logicalBranches = this._dynamicBranches;
     for (let i = 0, l = logicalBranches.length; i < l; i++) {

test/test_types.js

@@ -3505,6 +3505,54 @@ suite('types', () => {
       assert(Type.isType(t.field('unwrapped').type, 'union:unwrapped'));
     });
 
+    test('union projection', () => {
+      const Dog = {
+        type: 'record',
+        name: 'Dog',
+        fields: [
+          { type: 'string', name: 'bark' }
+        ],
+      };
+      const Cat = {
+        type: 'record',
+        name: 'Cat',
+        fields: [
+          { type: 'string', name: 'meow' }
+        ],
+      };
+      const animalTypes = [Dog, Cat];
+
+      const wrapUnions = (types) => {
+        assert.deepEqual(types.map(t => t.name), ['Dog', 'Cat']);
+        return (animal) => {
+          const animalType = ((animal) => {
+            if ('bark' in animal) {
+              return 'Dog';
+            } else if ('meow' in animal) {
+              return 'Cat';
+            }
+            throw new Error('Unknown animal');
+          })(animal);
+          return types.indexOf(types.find(type => type.name === animalType));
+        }
+      };
+
+      // TODO: replace this with a mock when available
+      // currently we're on mocha without sinon
+      function mockWrapUnions() {
+        mockWrapUnions.calls = typeof mockWrapUnions.calls === 'undefined'
+          ? 1
+          : ++mockWrapUnions.calls;
+        return wrapUnions.apply(null, arguments);
+      }
+
+       // Ambiguous, but we have a projection function
+      const Animal = Type.forSchema(animalTypes, { wrapUnions: mockWrapUnions });
+      Animal.toBuffer({ meow: '🐈' });
+      assert.equal(mockWrapUnions.calls, 1);
+      assert.throws(() => Animal.toBuffer({ snap: '🐊' }), /Unknown animal/)
+    });
+
     test('invalid wrap unions option', () => {
       assert.throws(() => {
         Type.forSchema('string', {wrapUnions: 'FOO'});

types/index.d.ts

@@ -96,6 +96,17 @@ interface EncoderOptions {
   syncMarker: Buffer;
 }
 
+/**
+ * A projection function that is used when unwrapping unions.
+ * This function is called at schema parsing time on each union with its branches' types.
+ * If it returns a non-null (function) value, that function will be called each time a value's branch needs to be inferred and should return the branch's index.
+ * The index muss be a number between 0 and length-1 of the passed types.
+ * In this case (a branch index) the union will use an unwrapped representation. Otherwise (undefined), the union will be wrapped.
+ */
+type BranchProjection = (types: ReadonlyArray<Type>) =>
+  | ((val: unknown) => number)
+  | undefined;
+
 interface ForSchemaOptions {
   assertLogicalTypes: boolean;
   logicalTypes: { [type: string]: new (schema: Schema, opts?: any) => types.LogicalType; };
@@ -104,7 +115,7 @@ interface ForSchemaOptions {
   omitRecordMethods: boolean;
   registry: { [name: string]: Type };
   typeHook: (schema: Schema | string, opts: ForSchemaOptions) => Type | undefined;
-  wrapUnions: boolean | 'auto' | 'always' | 'never';
+  wrapUnions: BranchProjection | boolean | 'auto' | 'always' | 'never';
 }
 
 interface TypeOptions extends ForSchemaOptions {