From f6afd22d3f5a20089759042f16fd865646a32038 Mon Sep 17 00:00:00 2001
From: Rikki Schulte <rikki.schulte@gmail.com>
Date: Tue, 31 Oct 2023 15:23:21 +0100
Subject: [PATCH] improve the state editor hooks, allow initialState

---
 .changeset/dirty-planes-count.md            |  5 ++
 packages/graphiql-react/src/editor/hooks.ts | 79 ++++++++++++---------
 2 files changed, 52 insertions(+), 32 deletions(-)
 create mode 100644 .changeset/dirty-planes-count.md

diff --git a/.changeset/dirty-planes-count.md b/.changeset/dirty-planes-count.md
new file mode 100644
index 00000000000..7724b10c970
--- /dev/null
+++ b/.changeset/dirty-planes-count.md
@@ -0,0 +1,5 @@
+---
+'@graphiql/react': minor
+---
+
+Add useHeadersEditorState and generic useEditorState hooks
diff --git a/packages/graphiql-react/src/editor/hooks.ts b/packages/graphiql-react/src/editor/hooks.ts
index a9c37f5da68..a7169ebc146 100644
--- a/packages/graphiql-react/src/editor/hooks.ts
+++ b/packages/graphiql-react/src/editor/hooks.ts
@@ -333,46 +333,61 @@ export function useAutoCompleteLeafs({
   }, [getDefaultFieldNames, queryEditor, schema]);
 }
 
+export type InitialState = string | (() => string);
+
 // https://react.dev/learn/you-might-not-need-an-effect
 
-/**
- * useState-like hook for current tab operations editor state
- */
-export function useOperationsEditorState(): [
-  opString: string,
-  handleEditOperations: (content: string) => void,
-] {
-  const { queryEditor } = useEditorContext({
+export const useEditorState = (
+  editor: 'query' | 'variable' | 'header',
+  initialState?: InitialState,
+) => {
+  const context = useEditorContext({
     nonNull: true,
   });
-  const opString = queryEditor?.getValue() ?? '';
-  const handleEditOperations = useCallback(
-    (value: string) => queryEditor?.setValue(value),
-    [queryEditor],
+  const initialValue =
+    typeof initialState === 'function' ? initialState() : initialState;
+  const editorInstance = context[`${editor}Editor` as const];
+  let valueString = '';
+  const editorValue = editorInstance?.getValue();
+  if (editorValue) {
+    valueString = editorValue;
+  } else {
+    valueString = initialValue || '';
+  }
+
+  const handleEditorValue = useCallback(
+    (value: string) => editorInstance?.setValue(value),
+    [editorInstance],
   );
-  return useMemo(
-    () => [opString, handleEditOperations],
-    [opString, handleEditOperations],
+  return useMemo<[string, (val: string) => void]>(
+    () => [valueString, handleEditorValue],
+    [valueString, handleEditorValue],
   );
+};
+
+/**
+ * useState-like hook for current tab operations editor state
+ */
+export function useOperationsEditorState(
+  initialState?: InitialState,
+): [operations: string, setOperations: (content: string) => void] {
+  return useEditorState('query', initialState);
 }
 
 /**
- * useState-like hook for variables tab operations editor state
+ * useState-like hook for current tab variables editor state
  */
-export function useVariablesEditorState(): [
-  varsString: string,
-  handleEditVariables: (content: string) => void,
-] {
-  const { variableEditor } = useEditorContext({
-    nonNull: true,
-  });
-  const varsString = variableEditor?.getValue() ?? '';
-  const handleEditVariables = useCallback(
-    (value: string) => variableEditor?.setValue(value),
-    [variableEditor],
-  );
-  return useMemo(
-    () => [varsString, handleEditVariables],
-    [varsString, handleEditVariables],
-  );
+export function useVariablesEditorState(
+  initialState?: InitialState,
+): [variables: string, setVariables: (content: string) => void] {
+  return useEditorState('variable', initialState);
+}
+
+/**
+ * useState-like hook for current tab variables editor state
+ */
+export function useHeadersEditorState(
+  initialState?: InitialState,
+): [headers: string, setHeaders: (content: string) => void] {
+  return useEditorState('header', initialState);
 }