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':
@@ -197,7 +199,20 @@ class Type {
         return Type.forSchema(obj, opts);
       });
       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 {
+          UnionType = typeof opts.wrapUnions(types) !== 'undefined'
+            // projection function yields a function, we can use an Unwrapped type
+            ? UnwrappedUnionType
+            : WrappedUnionType;
+          } catch(e) {
+            throw new Error(`Union projection function errored: ${e}`);
+          }
+        } else {
+          UnionType = isAmbiguous(types) ? WrappedUnionType : UnwrappedUnionType;
+        }
       }
       LOGICAL_TYPE = logicalType;
       type = new UnionType(types, opts);
@@ -341,10 +356,10 @@ class Type {
           return branchTypes[name];
         }), opts);
       } catch (err) {
-        opts.wrapUnions = wrapUnions;
         throw err;
+      } finally {
+        opts.wrapUnions = wrapUnions;
       }
-      opts.wrapUnions = wrapUnions;
       return unionType;
     }
 
@@ -1265,12 +1280,21 @@ class UnwrappedUnionType extends UnionType {
         this._bucketIndices[bucket] = index;
       }
     }, this);
+    this.projectionFunction = opts && typeof opts.wrapUnions === 'function'
+      ? ((fn) => (val) => {
+          const index = fn(val);
+          if (typeof index !== 'number' || index >= this._bucketIndices.length) {
+            throw new Error(`Projected index ${index} is not valid`);
+          }
+          return index;
+        })(opts.wrapUnions(this.types))
+      : (val) => this._bucketIndices[getValueBucket(val)]
 
     Object.freeze(this);
   }
 
   _getIndex (val) {
-    let index = this._bucketIndices[getValueBucket(val)];
+    let index = this.projectionFunction(val);
     if (this._dynamicBranches) {
       // Slower path, we must run the value through all branches.
       index = this._getBranchIndex(val, index);

test/test_types.js

@@ -3492,6 +3492,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, 2);
+      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,15 @@ 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 ProjectionFn = (types: ReadonlyArray<Type>) => ((val: unknown) => number) | undefined;
+
 interface ForSchemaOptions {
   assertLogicalTypes: boolean;
   logicalTypes: { [type: string]: new (schema: Schema, opts?: any) => types.LogicalType; };
@@ -104,7 +113,7 @@ interface ForSchemaOptions {
   omitRecordMethods: boolean;
   registry: { [name: string]: Type };
   typeHook: (schema: Schema | string, opts: ForSchemaOptions) => Type | undefined;
-  wrapUnions: boolean | 'auto' | 'always' | 'never';
+  wrapUnions: ProjectionFn | boolean | 'auto' | 'always' | 'never';
 }
 
 interface TypeOptions extends ForSchemaOptions {