From 54b091157fc3378a0caff9420acaa8f302535036 Mon Sep 17 00:00:00 2001
From: Oleh Misarosh <olehmisar@gmail.com>
Date: Fri, 17 Nov 2023 23:51:22 +0200
Subject: [PATCH] feat: typed context WIP

---
 packages/svelte/src/main/main-client.js | 15 +++-----
 packages/svelte/src/main/public.d.ts    | 20 ++++++++++
 packages/svelte/tests/types/context.ts  | 49 +++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 10 deletions(-)
 create mode 100644 packages/svelte/tests/types/context.ts

diff --git a/packages/svelte/src/main/main-client.js b/packages/svelte/src/main/main-client.js
index 0971ebdc14..a038f64d67 100644
--- a/packages/svelte/src/main/main-client.js
+++ b/packages/svelte/src/main/main-client.js
@@ -39,17 +39,11 @@ export function onMount(fn) {
 }
 
 /**
- * Retrieves the context that belongs to the closest parent component with the specified `key`.
- * Must be called during component initialisation.
- *
- * https://svelte.dev/docs/svelte#getcontext
- * @template T
  * @param {any} key
- * @returns {T}
  */
 export function getContext(key) {
 	const context_map = get_or_init_context_map();
-	return /** @type {T} */ (context_map.get(key));
+	return context_map.get(key);
 }
 
 /**
@@ -61,7 +55,7 @@ export function getContext(key) {
  *
  * https://svelte.dev/docs/svelte#setcontext
  * @template T
- * @param {any} key
+ * @param {{} | import('./public.js').ContextKey<T>} key
  * @param {T} context
  * @returns {T}
  */
@@ -76,8 +70,9 @@ export function setContext(key, context) {
  * Must be called during component initialisation.
  *
  * https://svelte.dev/docs/svelte#hascontext
- * @param {any} key
- * @returns {boolean}
+ * @template T
+ * @param {{} | import('./public.js').ContextKey<T>} key
+ * @returns {key is import('./public.js').CheckedContextKey<T>}
  */
 export function hasContext(key) {
 	const context_map = get_or_init_context_map();
diff --git a/packages/svelte/src/main/public.d.ts b/packages/svelte/src/main/public.d.ts
index f6d1ab7571..35622567e2 100644
--- a/packages/svelte/src/main/public.d.ts
+++ b/packages/svelte/src/main/public.d.ts
@@ -202,5 +202,25 @@ export interface EventDispatcher<EventMap extends Record<string, any>> {
 	): boolean;
 }
 
+/**
+ * Provided as key to `setContext`, `hasContext` and `getContext` in order to enable strict typing
+ */
+export interface ContextKey<T> extends Symbol {}
+
+declare const isChecked: unique symbol;
+export interface CheckedContextKey<T> extends ContextKey<T> {
+	[isChecked]: undefined;
+}
+
+/**
+ * Retrieves the context that belongs to the closest parent component with the specified `key`.
+ * Must be called during component initialisation.
+ *
+ * https://svelte.dev/docs/svelte#getcontext
+ */
+export function getContext<T>(key: CheckedContextKey<T>): T;
+export function getContext<T>(key: ContextKey<T>): T | undefined;
+export function getContext<T>(key: {}): T;
+
 export * from './main-client.js';
 import './ambient.js';
diff --git a/packages/svelte/tests/types/context.ts b/packages/svelte/tests/types/context.ts
new file mode 100644
index 0000000000..04c343dc2a
--- /dev/null
+++ b/packages/svelte/tests/types/context.ts
@@ -0,0 +1,49 @@
+import { getContext, hasContext, setContext, type ContextKey } from 'svelte';
+
+// non ContextKey still works
+setContext('string', 0);
+setContext(5, 0);
+setContext(true, 0);
+setContext({}, 0);
+setContext([], 0);
+setContext(Symbol(), 0);
+let resGet: number;
+resGet = getContext<number>('string');
+resGet = getContext<number>(5);
+resGet = getContext<number>(true);
+resGet = getContext<number>({});
+resGet = getContext<number>([]);
+
+// because `Symbol` is structurally the same as `ContextKey`, it returns optional type
+const resSymbol: number | undefined = getContext<number>(Symbol());
+
+let resHas: boolean;
+resHas = hasContext('string');
+resHas = hasContext(5);
+resHas = hasContext(true);
+resHas = hasContext({});
+resHas = hasContext([]);
+resHas = hasContext(Symbol());
+
+// ContextKey works
+const stringKey: ContextKey<string> = Symbol();
+setContext(stringKey, 'hello');
+// @ts-expect-error: wrong type of context
+setContext(stringKey, 1);
+const res1: string | undefined = getContext(stringKey);
+// @ts-expect-error: should be optional
+const res2: string = getContext(stringKey);
+if (hasContext(stringKey)) {
+	const resChecked: string = getContext(stringKey);
+}
+// @ts-expect-error: wrong type of variable
+const res3: number = getContext(stringKey);
+
+// TODO: generic takes over the ContextKey and function signature is: `getContext<number>(key: {})`
+// instead of`getContext<number>(key: ContextKey<string>)`
+// // @ts-expect-error: wrong type of generic
+// getContext<number>(stringKey);
+
+// TODO: fix this
+// // @ts-expect-error: wrong type of context generic
+// const numberKey: ContextKey<number> = stringKey;