✨ Better typings for oneof and frequency (#556)

Related to #546
dubzzz · Mar 2, 2020 · 80e6d7d · 80e6d7d
1 parent 4c03ea3
commit 80e6d7d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 11 deletions.
diff --git a/src/check/arbitrary/FrequencyArbitrary.ts b/src/check/arbitrary/FrequencyArbitrary.ts
@@ -33,18 +33,28 @@ class FrequencyArbitrary<T> extends Arbitrary<T> {
   }
 }
 
+/**
+ * Infer the type of the Arbitrary produced by oneof
+ * given the type of the source arbitraries
+ */
+type FrequencyArbitraryType<Ts extends WeightedArbitrary<unknown>[]> = {
+  [K in keyof Ts]: Ts[K] extends WeightedArbitrary<infer U> ? U : never
+}[number];
+
 /**
  * For one of the values generated by `...warbs` - the probability of selecting the ith warb is of `warb[i].weight / sum(warb[j].weight)`
  *
  * **WARNING**: It expects at least one (Arbitrary, weight)
  *
  * @param warbs (Arbitrary, weight)s that might be called to produce a value
  */
-function frequency<T>(...warbs: WeightedArbitrary<T>[]): Arbitrary<T> {
+function frequency<Ts extends WeightedArbitrary<unknown>[]>(...warbs: Ts): Arbitrary<FrequencyArbitraryType<Ts>> {
   if (warbs.length === 0) {
     throw new Error('fc.frequency expects at least one parameter');
   }
-  return new FrequencyArbitrary([...warbs]);
+  return new FrequencyArbitrary<FrequencyArbitraryType<Ts>>([...warbs] as WeightedArbitrary<
+    FrequencyArbitraryType<Ts>
+  >[]);
 }
 
 export { frequency };
diff --git a/src/check/arbitrary/OneOfArbitrary.ts b/src/check/arbitrary/OneOfArbitrary.ts
@@ -16,18 +16,26 @@ class OneOfArbitrary<T> extends Arbitrary<T> {
   }
 }
 
+/**
+ * Infer the type of the Arbitrary produced by oneof
+ * given the type of the source arbitraries
+ */
+type OneOfArbitraryType<Ts extends Arbitrary<unknown>[]> = {
+  [K in keyof Ts]: Ts[K] extends Arbitrary<infer U> ? U : never
+}[number];
+
 /**
  * For one of the values generated by `...arbs` - with all `...arbs` equiprobable
  *
  * **WARNING**: It expects at least one arbitrary
  *
  * @param arbs Arbitraries that might be called to produce a value
  */
-function oneof<T>(...arbs: Arbitrary<T>[]): Arbitrary<T> {
+function oneof<Ts extends Arbitrary<unknown>[]>(...arbs: Ts): Arbitrary<OneOfArbitraryType<Ts>> {
   if (arbs.length === 0) {
     throw new Error('fc.oneof expects at least one parameter');
   }
-  return new OneOfArbitrary([...arbs]);
+  return new OneOfArbitrary<OneOfArbitraryType<Ts>>([...arbs] as Arbitrary<OneOfArbitraryType<Ts>>[]);
 }
 
 export { oneof };
diff --git a/test/type/index.test-d.ts b/test/type/index.test-d.ts
@@ -60,21 +60,16 @@ expectError(fc.tuple(fc.nat(), ''));
 
 // oneof arbitrary
 expectType<fc.Arbitrary<string>>(fc.oneof(fc.string(), fc.fullUnicodeString()));
-expectType<fc.Arbitrary<number | string>>(fc.oneof(fc.string() as fc.Arbitrary<number | string>, fc.nat()));
-expectError(fc.oneof(fc.string(), fc.nat())); // TODO Typings should be improved
+expectType<fc.Arbitrary<string | number>>(fc.oneof(fc.string(), fc.nat()));
 expectError(fc.oneof(fc.string(), '1'));
 
 // frequency arbitrary
 expectType<fc.Arbitrary<string>>(
   fc.frequency({ arbitrary: fc.string(), weight: 1 }, { arbitrary: fc.fullUnicodeString(), weight: 1 })
 );
 expectType<fc.Arbitrary<number | string>>(
-  fc.frequency(
-    { arbitrary: fc.string() as fc.Arbitrary<number | string>, weight: 1 },
-    { arbitrary: fc.nat(), weight: 1 }
-  )
+  fc.frequency({ arbitrary: fc.string(), weight: 1 }, { arbitrary: fc.nat(), weight: 1 })
 );
-expectError(fc.frequency({ arbitrary: fc.string(), weight: 1 }, { arbitrary: fc.nat(), weight: 1 })); // TODO Typings should be improved
 expectError(fc.frequency({ arbitrary: fc.string(), weight: 1 }, { arbitrary: '1', weight: 1 }));
 
 // option arbitrary