From 46ea1a97483357a031ce5229e31d7de3c690e16a Mon Sep 17 00:00:00 2001
From: Tom Anderson <tom@ando.gq>
Date: Sat, 25 May 2024 22:05:33 +1000
Subject: [PATCH] allow for `.subscribe` to be overloaded

---
 .changes/overload-subscription.md |  5 +++++
 examples/counter/bindings.ts      |  3 ++-
 examples/counter/index.ts         |  6 ++++--
 packages/client/src/client.ts     | 24 ++++++++++++++++++------
 packages/client/src/stream.ts     | 16 ++++++++++------
 5 files changed, 39 insertions(+), 15 deletions(-)
 create mode 100644 .changes/overload-subscription.md

diff --git a/.changes/overload-subscription.md b/.changes/overload-subscription.md
new file mode 100644
index 0000000..3fd698c
--- /dev/null
+++ b/.changes/overload-subscription.md
@@ -0,0 +1,5 @@
+---
+"@qubit-rs/client": minor:feat
+---
+
+allow for subscription method to be overloaded if only `on_data` is required.
diff --git a/examples/counter/bindings.ts b/examples/counter/bindings.ts
index 5709f81..9675df4 100644
--- a/examples/counter/bindings.ts
+++ b/examples/counter/bindings.ts
@@ -1,5 +1,6 @@
+import type { Stream } from "@qubit-rs/client";
 export type Metadata = { param_a: string, param_b: number, param_c: boolean, more_metadata: Metadata | null, };
 export type Test = { a: number, b: boolean, };
 export type User = { name: string, email: string, age: number, metadata: Metadata, };
 
-export type Server = { version: () => Promise<string>, count: () => Promise<number>, array: () => Promise<Array<string>>, user: { someHandler: (_id: string) => Promise<User>, create: (name: string, email: string, age: number) => Promise<User>, list: () => Promise<Array<Test>>, asdf: () => Promise<null> } };
\ No newline at end of file
+export type Server = { version: () => Promise<string>, count: () => Promise<number>, countdown: (min: number, max: number) => Stream<number>, array: () => Promise<Array<string>>, user: { someHandler: (_id: string) => Promise<User>, create: (name: string, email: string, age: number) => Promise<User>, list: () => Promise<Array<Test>>, asdf: () => Promise<null> } };
\ No newline at end of file
diff --git a/examples/counter/index.ts b/examples/counter/index.ts
index 7a9dabd..6919ca2 100644
--- a/examples/counter/index.ts
+++ b/examples/counter/index.ts
@@ -4,10 +4,10 @@ import type { Server } from "./bindings.ts";
 const client = ws<Server>("ws://localhost:9944/rpc");
 
 client.version().then((version) => console.log({ version })).catch(console.error);
-client.user.get("test").then((user) => console.log(user)).catch(console.error);
+client.user.someHandler("test").then((user) => console.log(user)).catch(console.error);
 client.count().then((value) => console.log({ value })).catch(console.error);
 
-await client.countdown(1, 4).subscribe({
+client.countdown(1, 4).subscribe({
 	on_data: (data) => {
 		console.log("countdown: ", data);
 	},
@@ -16,3 +16,5 @@ await client.countdown(1, 4).subscribe({
 	}
 });
 
+client.countdown(1, 4)
+	.subscribe((n) => console.log("number is", n))
diff --git a/packages/client/src/client.ts b/packages/client/src/client.ts
index 32187c0..fa08372 100644
--- a/packages/client/src/client.ts
+++ b/packages/client/src/client.ts
@@ -34,16 +34,28 @@ export function build_client<Server>(client: Client): Server {
                         }
                     });
 
-                    const subscribe: StreamSubscriber<any> = ({ on_data, on_end, on_error }) => {
-                        function error(e: Error) {
-                            if (on_error) {
-                                on_error(e);
+                    const subscribe: StreamSubscriber<any> = (handler) => {
+                        let on_data = (_: unknown) => {};
+                        let on_error = (_: Error) => {};
+                        let on_end = () => {};
+
+                        if (typeof handler === "function") {
+                            on_data = handler;
+                        } else {
+                            if (handler?.on_data) {
+                                on_data = handler.on_data;
+                            }
+                            if (handler?.on_error) {
+                                on_error = handler.on_error;
+                            }
+                            if (handler?.on_end) {
+                                on_end = handler.on_end;
                             }
                         }
 
                         // Make sure the client can handle susbcriptions
                         if (!client.subscribe) {
-                            error(new Error("client does not support subscriptions"));
+                            on_error(new Error("client does not support subscriptions"));
                             return () => {};
                         }
                         const subscribe = client.subscribe;
@@ -68,7 +80,7 @@ export function build_client<Server>(client: Client): Server {
                             // Result should be a subscription ID
                             if (typeof subscription_id !== "string" && typeof subscription_id !== "number") {
                                 // TODO: Throw an error
-                                error(new Error("cannot subscribe to subscription"));
+                                on_error(new Error("cannot subscribe to subscription"));
                                 return () => {};
                             }
 
diff --git a/packages/client/src/stream.ts b/packages/client/src/stream.ts
index 50a2782..929a083 100644
--- a/packages/client/src/stream.ts
+++ b/packages/client/src/stream.ts
@@ -1,9 +1,13 @@
-export type StreamSubscriber<T> = ({ on_data, on_error, on_end }: {
-	on_data?: (data: T) => void,
-	on_error?: (error: Error) => void,
-	on_end?: () => void,
-}) => () => void;
+export type StreamSubscriber<T> = (
+  handler:
+    | ((data: T) => void)
+    | {
+        on_data?: (data: T) => void;
+        on_error?: (error: Error) => void;
+        on_end?: () => void;
+      },
+) => () => void;
 
 export type Stream<T> = {
-	subscribe: StreamSubscriber<T>,
+  subscribe: StreamSubscriber<T>;
 };