update subscribe with middleware

immutable/package.json

@@ -1,6 +1,6 @@
 {
   "name": "swr-immutable",
-  "version": "0.0.1",
+  "private": true,
   "main": "./dist/index.js",
   "module": "./dist/index.esm.js",
   "types": "./dist/immutable",

infinite/package.json

@@ -1,6 +1,6 @@
 {
   "name": "swr-infinite",
-  "version": "0.0.1",
+  "private": true,
   "main": "./dist/index.js",
   "module": "./dist/index.esm.js",
   "types": "./dist/infinite",

package.json

@@ -34,12 +34,14 @@
   "homepage": "https://swr.vercel.app",
   "license": "MIT",
   "scripts": {
-    "build": "yarn build:core && yarn build:infinite && yarn build:immutable",
+    "clean": "rm -rf dist infinite/dist immutable/dist subscribe/dist",
+    "build": "yarn build:core && yarn build:infinite && yarn build:immutable && yarn build:subscribe",
     "watch": "bunchee src/index.ts --watch",
     "build:core": "bunchee src/index.ts -m --no-sourcemap",
     "build:infinite": "bunchee index.ts --cwd infinite -m --no-sourcemap",
     "build:immutable": "bunchee index.ts --cwd immutable -m --no-sourcemap",
-    "prepublishOnly": "rm -rf dist infinite/dist immutable/dist && yarn build",
+    "build:subscribe": "bunchee index.ts --cwd subscribe -m --no-sourcemap",
+    "prepublishOnly": "yarn clean && yarn build",
     "types:check": "tsc --noEmit",
     "format": "prettier --write \"{src,infinite,immutable,test,examples}/**/*.{ts,tsx}\"",
     "lint": "eslint \"{src,infinite,immutable,test,examples}/**/*.{ts,tsx}\"",

subscribe/index.ts

@@ -0,0 +1,61 @@
+import { useEffect } from 'react'
+
+// @ts-expect-error
+import useSWR from 'swr'
+import { withMiddleware } from '../src/utils/with-middleware'
+import { Key, SWRHook, Middleware, SWRConfiguration } from '../src/types'
+
+export type SWRSubscription<Data, Error> = (
+  key: Key,
+  callbacks: {
+    onData(data: Data): void
+    onError(err: Error): void
+  }
+) => void
+
+export type SWRSubscriptionResponse<Data, Error> = {
+  data?: Data
+  error?: Error
+}
+
+const subscriptions = new Map<Key, number>()
+const disposers = new Map()
+
+export const subscribe = (<Data, Error>(useSWRNext: SWRHook) => (
+  key: Key,
+  subscribeFn: SWRSubscription<Data, Error>,
+  config?: SWRConfiguration
+): SWRSubscriptionResponse<Data, Error> => {
+  const { data, error, mutate } = useSWRNext(key, null, config)
+
+  useEffect(() => {
+    subscriptions.set(key, (subscriptions.get(key) || 0) + 1)
+    const onData = (val: Data) => mutate(val, false)
+    const onError = async (err: any) => {
+      // Avoid thrown errors from `mutate`
+      // eslint-disable-next-line no-empty
+      try {
+        await mutate(() => {
+          throw err
+        }, false)
+      } catch (_) {}
+    }
+
+    if (subscriptions.get(key) === 1) {
+      const dispose = subscribeFn(key, { onData, onError })
+      disposers.set(key, dispose)
+    }
+    return () => {
+      const count = subscriptions.get(key) || 1
+      subscriptions.set(key, count - 1)
+      // dispose if it's last one
+      if (count === 1) {
+        disposers.get(key)()
+      }
+    }
+  }, [key])
+
+  return { data, error }
+}) as Middleware
+
+export default withMiddleware(useSWR, subscribe)

subscribe/package.json

@@ -0,0 +1,12 @@
+{
+  "name": "swr-subscribe",
+  "private": true,
+  "main": "./dist/index.js",
+  "module": "./dist/index.esm.js",
+  "types": "./dist/subscribe",
+  "peerDependencies": {
+    "swr": "*",
+    "react": "*",
+    "dequal": "*"
+  }
+}

subscribe/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "./dist",
+    "rootDir": "..",
+  },
+  "include": [".", "../src"],
+  "exclude": ["./dist"]
+}

test/use-swr-subscribe.test.tsx

@@ -0,0 +1,49 @@
+import React from 'react'
+import { render } from '@testing-library/react'
+import { sleep } from './utils'
+import useSWRSubscription from '../subscribe'
+
+describe('useSWRSubscription', () => {
+  it('should update state when fetcher is a subscription', async () => {
+    const key = 'sub-0'
+    let intervalId
+    let res = 0
+    function subscribe(_key, { onData, onError }) {
+      intervalId = setInterval(() => {
+        if (res === 3) {
+          const err = new Error(_key + 'error')
+          onError(err)
+        } else {
+          onData(_key + res)
+        }
+        res++
+      }, 100)
+
+      return () => {}
+    }
+
+    function Page() {
+      const { data, error } = useSWRSubscription(key, subscribe)
+      return <div>{error ? error.message : data}</div>
+    }
+
+    const { container } = render(<Page />)
+    await sleep(10)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(`""`)
+    await sleep(100)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(`"${key}0"`)
+    await sleep(100)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(`"${key}1"`)
+    await sleep(100)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(`"${key}2"`)
+    await sleep(100)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(
+      `"${key}error"`
+    )
+    clearInterval(intervalId)
+    await sleep(100)
+    expect(container.firstChild.textContent).toMatchInlineSnapshot(
+      `"${key}error"`
+    )
+  })
+})