Added selection-aware word count plugin for Koenig

packages/koenig-lexical/demo/DemoApp.jsx

@@ -17,7 +17,7 @@ import {$getRoot, $isDecoratorNode} from 'lexical';
 import {
     BASIC_NODES, BASIC_TRANSFORMERS, EmailEditor,
     KoenigComposableEditor, KoenigComposer, KoenigEditor, MINIMAL_NODES,
-    MINIMAL_TRANSFORMERS, RestrictContentPlugin, TKCountPlugin, WordCountPlugin
+    MINIMAL_TRANSFORMERS, RestrictContentPlugin, SelectionWordCountPlugin, TKCountPlugin, WordCountPlugin
 } from '../src';
 import {defaultHeaders as defaultUnsplashHeaders} from './utils/unsplashConfig';
 import {fetchEmbed} from './utils/fetchEmbed';
@@ -131,7 +131,7 @@ function getAllowedNodes({editorType}) {
     return undefined;
 }
 
-function DemoEditor({editorType, registerAPI, cursorDidExitAtTop, darkMode, setWordCount, setTKCount}) {
+function DemoEditor({editorType, registerAPI, cursorDidExitAtTop, darkMode, setWordCount, setSelectionWordCount, setTKCount}) {
     if (editorType === 'basic') {
         return (
             <KoenigComposableEditor
@@ -140,6 +140,7 @@ function DemoEditor({editorType, registerAPI, cursorDidExitAtTop, darkMode, setW
                 registerAPI={registerAPI}
             >
                 <WordCountPlugin onChange={setWordCount} />
+                <SelectionWordCountPlugin onChange={setSelectionWordCount} />
             </KoenigComposableEditor>
         );
     } else if (editorType === 'minimal') {
@@ -152,6 +153,7 @@ function DemoEditor({editorType, registerAPI, cursorDidExitAtTop, darkMode, setW
             >
                 <RestrictContentPlugin paragraphs={1} />
                 <WordCountPlugin onChange={setWordCount} />
+                <SelectionWordCountPlugin onChange={setSelectionWordCount} />
             </KoenigComposableEditor>
         );
     }
@@ -163,12 +165,13 @@ function DemoEditor({editorType, registerAPI, cursorDidExitAtTop, darkMode, setW
             registerAPI={registerAPI}
         >
             <WordCountPlugin onChange={setWordCount} />
+            <SelectionWordCountPlugin onChange={setSelectionWordCount} />
             <TKCountPlugin onChange={setTKCount} />
         </KoenigEditor>
     );
 }
 
-function DemoComposer({editorType, isMultiplayer, setWordCount, setTKCount}) {
+function DemoComposer({editorType, isMultiplayer, setWordCount, setSelectionWordCount, setTKCount}) {
     const [searchParams, setSearchParams] = useSearchParams();
     const [isSidebarOpen, setIsSidebarOpen] = useState(false);
     const [sidebarView, setSidebarView] = useState('json');
@@ -405,6 +408,7 @@ function DemoComposer({editorType, isMultiplayer, setWordCount, setTKCount}) {
                     darkMode={darkMode}
                     editorType={editorType}
                     registerAPI={setEditorAPI}
+                    setSelectionWordCount={setSelectionWordCount}
                     setTKCount={setTKCount}
                     setWordCount={setWordCount}
                 />
@@ -418,6 +422,7 @@ const MemoizedDemoComposer = React.memo(DemoComposer);
 
 function DemoApp({editorType, isMultiplayer}) {
     const [wordCount, setWordCount] = useState(0);
+    const [selectionWordCount, setSelectionWordCount] = useState(null);
     const [tkCount, setTKCount] = useState(0);
 
     // used to force a re-initialization of the editor when URL changes, otherwise
@@ -430,11 +435,12 @@ function DemoApp({editorType, isMultiplayer}) {
             className={`koenig-lexical top`}
         >
             {/* outside of DemoComposer to avoid re-renders and flaky tests when word count changes */}
-            <WordCount tkCount={tkCount} wordCount={wordCount} />
+            <WordCount selectionWordCount={selectionWordCount} tkCount={tkCount} wordCount={wordCount} />
 
             <MemoizedDemoComposer
                 editorType={editorType}
                 isMultiplayer={isMultiplayer}
+                setSelectionWordCount={setSelectionWordCount}
                 setTKCount={setTKCount}
                 setWordCount={setWordCount}
             />

packages/koenig-lexical/demo/components/WordCount.jsx

@@ -1,6 +1,12 @@
-const WordCount = ({wordCount, tkCount}) => {
+const WordCount = ({wordCount, selectionWordCount, tkCount}) => {
     return (
         <div className="absolute left-6 top-4 z-20 block cursor-pointer rounded bg-white px-2 py-1 font-mono text-sm tracking-tight text-grey-600 dark:bg-transparent">
+            {selectionWordCount !== null && (
+                <>
+                    <span data-testid="selection-word-count">{selectionWordCount}</span>
+                    {' of '}
+                </>
+            )}
             <span data-testid="word-count">{wordCount}</span> words
             {tkCount > 0 && (
                 <>

packages/koenig-lexical/src/components/KoenigComposer.jsx

@@ -70,6 +70,11 @@ const KoenigComposer = ({
 
     const editorContainerRef = React.useRef(null);
     const onWordCountChangeRef = React.useRef(null);
+    const onSelectionWordCountChangeRef = React.useRef(null);
+    const selectionWordCountsRef = React.useRef(new Map());
+    // starts as undefined (not null) so SelectionWordCountPlugin's first
+    // emission always fires, including the initial no-selection null
+    const lastEmittedSelectionWordCountRef = React.useRef(undefined);
 
     if (!fileUploader.useFileUpload) {
         fileUploader.useFileUpload = function () {
@@ -118,7 +123,10 @@ const KoenigComposer = ({
                 multiplayerDocId,
                 multiplayerUsername,
                 createWebsocketProvider,
-                onWordCountChangeRef
+                onWordCountChangeRef,
+                onSelectionWordCountChangeRef,
+                selectionWordCountsRef,
+                lastEmittedSelectionWordCountRef
             }}>
                 <KoenigSelectedCardContext>
                     <TKContext>

packages/koenig-lexical/src/components/KoenigNestedComposer.jsx

@@ -1,13 +1,14 @@
 import KoenigComposerContext from '../context/KoenigComposerContext';
 import React from 'react';
+import SelectionWordCountPlugin from '../plugins/SelectionWordCountPlugin';
 import WordCountPlugin from '../plugins/WordCountPlugin';
 import {CollaborationPlugin} from '@lexical/react/LexicalCollaborationPlugin';
 import {LexicalNestedComposer} from '@lexical/react/LexicalNestedComposer';
 import {useCollaborationContext} from '@lexical/react/LexicalCollaborationContext';
 
 const KoenigNestedComposer = ({initialEditor, initialEditorState, initialNodes, initialTheme, skipCollabChecks, children} = {}) => {
     const {isCollabActive} = useCollaborationContext();
-    const {createWebsocketProvider, onWordCountChangeRef} = React.useContext(KoenigComposerContext);
+    const {createWebsocketProvider, onWordCountChangeRef, onSelectionWordCountChangeRef} = React.useContext(KoenigComposerContext);
 
     return (
         <LexicalNestedComposer
@@ -27,6 +28,9 @@ const KoenigNestedComposer = ({initialEditor, initialEditorState, initialNodes,
             {onWordCountChangeRef?.current ? (
                 <WordCountPlugin onChange={onWordCountChangeRef.current} />
             ) : null}
+            {onSelectionWordCountChangeRef?.current ? (
+                <SelectionWordCountPlugin onChange={onSelectionWordCountChangeRef.current} />
+            ) : null}
             {children}
         </LexicalNestedComposer>
     );

packages/koenig-lexical/src/index.js

@@ -39,6 +39,7 @@ import PlusCardMenuPlugin from './plugins/PlusCardMenuPlugin';
 import ProductPlugin from './plugins/ProductPlugin';
 import ReplacementStringsPlugin from './plugins/ReplacementStringsPlugin';
 import RestrictContentPlugin from './plugins/RestrictContentPlugin';
+import SelectionWordCountPlugin from './plugins/SelectionWordCountPlugin';
 import SignupPlugin from './plugins/SignupPlugin';
 import SlashCardMenuPlugin from './plugins/SlashCardMenuPlugin';
 import TKCountPlugin from './plugins/TKCountPlugin';
@@ -112,6 +113,7 @@ export {
     ProductPlugin,
     ReplacementStringsPlugin,
     RestrictContentPlugin,
+    SelectionWordCountPlugin,
     SignupPlugin,
     SlashCardMenuPlugin,
     TKCountPlugin,

packages/koenig-lexical/src/plugins/SelectionWordCountPlugin.jsx

@@ -0,0 +1,92 @@
+import KoenigComposerContext from '../context/KoenigComposerContext';
+import React from 'react';
+import throttle from 'lodash/throttle';
+import {$getSelection, $isRangeSelection} from 'lexical';
+import {useLexicalComposerContext} from '@lexical/react/LexicalComposerContext';
+import {utils} from '@tryghost/helpers';
+
+const {countWords} = utils;
+
+// Reports the selection's word count: an integer (0 allowed) while a
+// non-collapsed range selection exists, null otherwise.
+export const SelectionWordCountPlugin = ({onChange} = {}) => {
+    const [editor] = useLexicalComposerContext();
+    const {onSelectionWordCountChangeRef, selectionWordCountsRef, lastEmittedSelectionWordCountRef} = React.useContext(KoenigComposerContext);
+
+    React.useLayoutEffect(() => {
+        if (!onChange) {
+            return;
+        }
+
+        // store onChange in context so that KoenigNestedComposer can render
+        // a nested <SelectionWordCountPlugin /> without passing onChange down
+        if (!editor._parentEditor) {
+            onSelectionWordCountChangeRef.current = onChange;
+        }
+
+        // a selection lives in exactly one editor at a time but a plugin
+        // instance exists per editor (main + nested caption editors). Each
+        // instance writes its own editor's selection count into the shared
+        // map and emits the combined value so that the emission order
+        // between editor instances doesn't matter
+        const counts = selectionWordCountsRef.current;
+
+        const emitCombined = () => {
+            let combined = null;
+
+            for (const count of counts.values()) {
+                if (count !== null) {
+                    combined = count;
+                    break;
+                }
+            }
+
+            // dedup lives on the composer-level ref rather than in this
+            // closure so that no instance re-emits a value another
+            // instance has already delivered
+            if (combined !== lastEmittedSelectionWordCountRef.current) {
+                lastEmittedSelectionWordCountRef.current = combined;
+                onChange(combined);
+            }
+        };
+
+        const updateSelectionCount = () => {
+            editor.getEditorState().read(() => {
+                const selection = $getSelection();
+
+                if ($isRangeSelection(selection) && !selection.isCollapsed()) {
+                    counts.set(editor.getKey(), countWords(selection.getTextContent()));
+                } else {
+                    counts.set(editor.getKey(), null);
+                }
+            });
+
+            emitCombined();
+        };
+
+        updateSelectionCount();
+
+        const throttledUpdate = throttle(updateSelectionCount, 200);
+
+        // update listeners fire after every commit, including selection-only
+        // changes, so a SELECTION_CHANGE_COMMAND handler isn't needed
+        const cleanupRegister = editor.registerUpdateListener(() => {
+            throttledUpdate();
+        });
+
+        return () => {
+            throttledUpdate.cancel();
+            cleanupRegister();
+            counts.delete(editor.getKey());
+            // an unmounting editor may hold the active selection (e.g. a
+            // card deleted mid-edit) so re-emit to drop its count
+            emitCombined();
+
+            if (!editor._parentEditor) {
+                onSelectionWordCountChangeRef.current = null;
+            }
+        };
+    }, [editor, onChange, onSelectionWordCountChangeRef, selectionWordCountsRef, lastEmittedSelectionWordCountRef]);
+};
+
+export default SelectionWordCountPlugin;

packages/koenig-lexical/test/e2e/plugins/SelectionWordCountPlugin.test.js

@@ -0,0 +1,118 @@
+import {ctrlOrCmd, focusEditor, initialize, insertCard, selectBackwards} from '../../utils/e2e';
+import {expect, test} from '@playwright/test';
+
+test.describe('Selection Word Count Plugin', async function () {
+    let page;
+
+    test.beforeAll(async ({browser}) => {
+        page = await browser.newPage();
+    });
+
+    test.beforeEach(async () => {
+        await initialize({page});
+    });
+
+    test.afterAll(async () => {
+        await page.close();
+    });
+
+    test('shows no selection count without a selection', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello beautiful world');
+        await expect(page.getByTestId('word-count')).toHaveText('3');
+        await expect(page.getByTestId('selection-word-count')).not.toBeVisible();
+    });
+
+    test('shows selection count for selected text', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello beautiful world');
+        await selectBackwards(page, 5); // selects "world"
+        await expect(page.getByTestId('selection-word-count')).toHaveText('1');
+        // total stays visible alongside the selection count
+        await expect(page.getByTestId('word-count')).toHaveText('3');
+
+        // extend across a word boundary; selects "beautiful world"
+        await selectBackwards(page, 10);
+        await expect(page.getByTestId('selection-word-count')).toHaveText('2');
+    });
+
+    test('shows 0 for a whitespace-only selection', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        for (let i = 0; i < 5; i++) {
+            await page.keyboard.press('ArrowLeft');
+        }
+
+        // Chrome for Testing occasionally drops the programmatic selection;
+        // retry until the whitespace-only selection registers and gets counted
+        await expect(async () => {
+            if (!(await page.evaluate(() => window.getSelection().isCollapsed))) {
+                // a previous attempt selected; collapse back to just before "world"
+                await page.keyboard.press('ArrowRight');
+            }
+            await selectBackwards(page, 1); // selects the space between the words
+            await expect(page.getByTestId('selection-word-count')).toHaveText('0', {timeout: 1000});
+        }).toPass();
+    });
+
+    test('counts partial word fragments as words', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        await selectBackwards(page, 3); // selects "rld"
+        await expect(page.getByTestId('selection-word-count')).toHaveText('1');
+    });
+
+    test('reverts to total-only display when selection collapses', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        await selectBackwards(page, 5);
+        await expect(page.getByTestId('selection-word-count')).toHaveText('1');
+        await page.keyboard.press('ArrowRight'); // collapse the selection
+        await expect(page.getByTestId('selection-word-count')).not.toBeVisible();
+        await expect(page.getByTestId('word-count')).toHaveText('2');
+    });
+
+    test('counts selection inside a nested editor', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        await page.keyboard.press('Enter');
+        await insertCard(page, {cardName: 'callout'});
+        await page.keyboard.type('Nested content here');
+        await expect(page.getByTestId('word-count')).toHaveText('5');
+        await selectBackwards(page, 12); // selects "content here"
+        await expect(page.getByTestId('selection-word-count')).toHaveText('2');
+    });
+
+    test('moving selection from nested editor back to main editor reports main selection', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        await page.keyboard.press('Enter');
+        await insertCard(page, {cardName: 'callout'});
+        await page.keyboard.type('Nested content here');
+        await selectBackwards(page, 4); // selects "here" in the nested editor
+        await expect(page.getByTestId('selection-word-count')).toHaveText('1');
+
+        // click into the first paragraph and select "Hello world" there
+        await page.locator('[data-lexical-editor] > p').first().click();
+        await page.keyboard.press('End');
+        await selectBackwards(page, 11);
+        await expect(page.getByTestId('selection-word-count')).toHaveText('2');
+
+        // collapse — counter must revert to total-only, with no stale nested count
+        await page.keyboard.press('ArrowRight');
+        await expect(page.getByTestId('selection-word-count')).not.toBeVisible();
+    });
+
+    test('select-all across a card matches the document total', async function () {
+        await focusEditor(page);
+        await page.keyboard.type('Hello world');
+        await page.keyboard.press('Enter');
+        await insertCard(page, {cardName: 'callout'});
+        await page.keyboard.type('Nested content here');
+        await page.keyboard.press('Escape'); // exit nested editing, card selected
+        await page.keyboard.press('ArrowDown'); // cursor into main editor below card
+        await page.keyboard.press(`${ctrlOrCmd(page)}+KeyA`);
+        await expect(page.getByTestId('word-count')).toHaveText('5');
+        await expect(page.getByTestId('selection-word-count')).toHaveText('5');
+    });
+});