From ccb2065d1147af56b5c1a929635f63443e469205 Mon Sep 17 00:00:00 2001
From: Alex Potsides <alex@achingbrain.net>
Date: Wed, 1 Mar 2023 15:18:40 +0000
Subject: [PATCH] fix: make pbStream function polymorphic (#40)

What we want to say is that the type returned from `unwrap` will be
the original duplex passed in, but it's source will have changed to
return `Uint8ArrayList`s.  We can still pass duplexes that source
`Uint8Array`s though, so make the exported function polymorphic.
---
 src/index.ts       | 20 ++++++++++++--------
 test/index.spec.ts |  2 +-
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 503cfb6..f6991a5 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -43,7 +43,7 @@ export interface Encoder<T> {
 /**
  * Convenience methods for working with protobuf streams
  */
-export interface ProtobufStream <TSink> {
+export interface ProtobufStream <Stream extends Duplex<Uint8ArrayList, Uint8ArrayList | Uint8Array> = Duplex<Uint8ArrayList, Uint8ArrayList | Uint8Array>> {
   /**
    * Read a set number of bytes from the stream
    */
@@ -82,7 +82,7 @@ export interface ProtobufStream <TSink> {
   /**
    * Returns the underlying stream
    */
-  unwrap: () => Duplex<Uint8ArrayList, TSink>
+  unwrap: () => Stream
 }
 
 export interface Opts {
@@ -97,14 +97,16 @@ export interface Opts {
   maxDataLength: number
 }
 
-export function pbStream <TSink extends Uint8Array | Uint8ArrayList> (duplex: Duplex<Uint8ArrayList | Uint8Array, TSink>, opts = {}): ProtobufStream<TSink> {
-  const shake = handshake(duplex)
+export function pbStream <Stream extends Duplex<Uint8ArrayList, Uint8Array | Uint8ArrayList>> (duplex: Stream, opts?: Partial<Opts>): ProtobufStream<Stream>
+export function pbStream <Stream extends Duplex<Uint8ArrayList, Uint8Array | Uint8ArrayList>> (duplex: Duplex<Uint8Array>, opts?: Partial<Opts>): ProtobufStream<Stream>
+export function pbStream (duplex: any, opts = {}): ProtobufStream<any> {
+  const shake = handshake<Uint8Array>(duplex)
   const lpReader = lp.decode.fromReader(
     shake.reader,
     opts
   )
 
-  const W: ProtobufStream<TSink> = {
+  const W: ProtobufStream<any> = {
     read: async (bytes) => {
       // just read
       const { value } = await shake.reader.next(bytes)
@@ -142,10 +144,8 @@ export function pbStream <TSink extends Uint8Array | Uint8ArrayList> (duplex: Du
     write: (data) => {
       // just write
       if (data instanceof Uint8Array) {
-        // @ts-expect-error writer should always accept pushing Uint8Arrays
         shake.writer.push(data)
       } else {
-        // @ts-expect-error writer should always accept pushing Uint8Arrays
         shake.writer.push(data.subarray())
       }
     },
@@ -166,7 +166,11 @@ export function pbStream <TSink extends Uint8Array | Uint8ArrayList> (duplex: Du
     unwrap: () => {
       // returns vanilla duplex again, terminates all reads/writes from this object
       shake.rest()
-      return shake.stream
+
+      duplex.source = shake.stream.source
+      duplex.sink = shake.stream.sink
+
+      return duplex
     }
   }
 
diff --git a/test/index.spec.ts b/test/index.spec.ts
index 7f66d02..b747a52 100644
--- a/test/index.spec.ts
+++ b/test/index.spec.ts
@@ -65,7 +65,7 @@ Object.keys(tests).forEach(key => {
   const test = tests[key]
 
   describe(`it-pb-rpc ${key}`, () => {
-    let w: ProtobufStream<Uint8Array>
+    let w: ProtobufStream
 
     before(async () => {
       w = pbStream(pair<Uint8Array>())