diff --git a/src/Condition.js b/src/Condition.js index 60a147e5e..a40d0f16c 100644 --- a/src/Condition.js +++ b/src/Condition.js @@ -1,6 +1,13 @@ import has from 'lodash/has'; import isSchema from './util/isSchema'; +function callOrConcat(schema) { + if (typeof schema === 'function') + return schema + + return base => base.concat(schema) +} + class Conditional { constructor(refs, options) { @@ -8,6 +15,9 @@ class Conditional { this.refs = [].concat(refs) + then = callOrConcat(then); + otherwise = callOrConcat(otherwise); + if (typeof options === 'function') this.fn = options else @@ -24,8 +34,10 @@ class Conditional { ? is : ((...values) => values.every(value => value === is)) this.fn = function (...values) { - let ctx = values.pop(); - return isFn(...values) ? ctx.concat(then) : ctx.concat(otherwise) + let currentSchema = values.pop(); + let option = isFn(...values) ? then : otherwise + + return option(currentSchema) } } }