From 80e6d7d8f6d489d3b56e2a044b54bdacdaea6553 Mon Sep 17 00:00:00 2001
From: Nicolas DUBIEN <github@dubien.org>
Date: Mon, 2 Mar 2020 23:51:10 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Better=20typings=20for=20oneof=20an?=
 =?UTF-8?q?d=20frequency=20(#556)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Related to #546
---
 src/check/arbitrary/FrequencyArbitrary.ts | 14 ++++++++++++--
 src/check/arbitrary/OneOfArbitrary.ts     | 12 ++++++++++--
 test/type/index.test-d.ts                 |  9 ++-------
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/check/arbitrary/FrequencyArbitrary.ts b/src/check/arbitrary/FrequencyArbitrary.ts
index 0f7567a7a32..19df15426d3 100644
--- a/src/check/arbitrary/FrequencyArbitrary.ts
+++ b/src/check/arbitrary/FrequencyArbitrary.ts
@@ -33,6 +33,14 @@ 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)`
  *
@@ -40,11 +48,13 @@ class FrequencyArbitrary<T> extends Arbitrary<T> {
  *
  * @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
index fa14f74877b..2979ef56b6e 100644
--- a/src/check/arbitrary/OneOfArbitrary.ts
+++ b/src/check/arbitrary/OneOfArbitrary.ts
@@ -16,6 +16,14 @@ 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
  *
@@ -23,11 +31,11 @@ class OneOfArbitrary<T> extends Arbitrary<T> {
  *
  * @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
index b3b7e1fcea7..8cd436fd636 100644
--- a/test/type/index.test-d.ts
+++ b/test/type/index.test-d.ts
@@ -60,8 +60,7 @@ 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
@@ -69,12 +68,8 @@ 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