refactor

microsoft · Nov 21, 2024 · a91d593 · a91d593
1 parent 7efa7de
commit a91d593
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 50 deletions.
diff --git a/packages/roosterjs-content-model-api/lib/modelApi/common/queryContentModelBlocks.ts b/packages/roosterjs-content-model-api/lib/modelApi/common/queryContentModelBlocks.ts
@@ -1,75 +1,61 @@
 import type {
     ContentModelBlockType,
     ReadonlyContentModelBlock,
+    ReadonlyContentModelBlockBase,
     ReadonlyContentModelBlockGroup,
 } from 'roosterjs-content-model-types';
 
 /**
  * Query content model blocks
  * @param group The block group to query
- * @param blockType The type of block to query @default 'Paragraph'
+ * @param type The type of block to query
  * @param filter Optional selector to filter the blocks
  * @param findFirstOnly True to return the first block only, false to return all blocks
  */
 export function queryContentModelBlocks<T extends ReadonlyContentModelBlock>(
     group: ReadonlyContentModelBlockGroup,
-    blockType?: ContentModelBlockType,
+    type: T extends ReadonlyContentModelBlockBase<infer U> ? U : never,
     filter?: (element: T) => element is T,
     findFirstOnly?: boolean
 ): T[] {
-    const type = blockType || 'Paragraph';
     const elements: T[] = [];
     for (let i = 0; i < group.blocks.length; i++) {
         if (findFirstOnly && elements.length > 0) {
             return elements;
         }
         const block = group.blocks[i];
-        const results = queryContentModelBlocksInternal<T>(block, type, filter, findFirstOnly);
-        elements.push(...results);
-    }
-    return elements;
-}
 
-function queryContentModelBlocksInternal<T extends ReadonlyContentModelBlock>(
-    block: ReadonlyContentModelBlock,
-    type: ContentModelBlockType,
-    filter?: (element: T) => element is T,
-    findFirstOnly?: boolean
-): T[] {
-    const elements: T[] = [];
-    if (isExpectedBlockType(block, type, filter)) {
-        elements.push(block);
-    }
-
-    if (block.blockType == 'BlockGroup') {
-        for (const childBlock of block.blocks) {
-            if (findFirstOnly && elements.length > 0) {
-                return elements;
-            }
-            const results = queryContentModelBlocksInternal<T>(
-                childBlock,
-                type,
-                filter,
-                findFirstOnly
-            );
-            elements.push(...results);
-        }
-    }
-
-    if (block.blockType == 'Table') {
-        const table = block;
-        for (const row of table.rows) {
-            for (const cell of row.cells) {
-                for (const cellBlock of cell.blocks) {
-                    const results = queryContentModelBlocksInternal<T>(
-                        cellBlock,
-                        type,
-                        filter,
-                        findFirstOnly
-                    );
-                    elements.push(...results);
+        switch (block.blockType) {
+            case 'Paragraph':
+            case 'Divider':
+            case 'Entity':
+                if (isExpectedBlockType(block, type, filter)) {
+                    elements.push(block);
+                }
+                break;
+            case 'BlockGroup':
+                if (isExpectedBlockType(block, type, filter)) {
+                    elements.push(block);
+                }
+                const results = queryContentModelBlocks<T>(block, type, filter, findFirstOnly);
+                elements.push(...results);
+                break;
+            case 'Table':
+                if (isExpectedBlockType(block, type, filter)) {
+                    elements.push(block);
+                }
+                for (const row of block.rows) {
+                    for (const cell of row.cells) {
+                        const results = queryContentModelBlocks<T>(
+                            cell,
+                            type,
+                            filter,
+                            findFirstOnly
+                        );
+                        elements.push(...results);
+                    }
                 }
-            }
+                break;
         }
     }
     return elements;
@@ -85,7 +71,7 @@ function isExpectedBlockType<T extends ReadonlyContentModelBlock>(
 
 function isBlockType<T extends ReadonlyContentModelBlock>(
     block: ReadonlyContentModelBlock,
-    type: string
+    type: ContentModelBlockType
 ): block is T {
     return block.blockType == type;
 }
diff --git a/packages/roosterjs-content-model-api/test/modelApi/common/queryContentModelBlocksTest.ts b/packages/roosterjs-content-model-api/test/modelApi/common/queryContentModelBlocksTest.ts
@@ -15,7 +15,7 @@ describe('queryContentModelBlocksBlocks', () => {
         };
 
         // Act
-        const result = queryContentModelBlocks(group);
+        const result = queryContentModelBlocks<ReadonlyContentModelParagraph>(group, 'Paragraph');
 
         // Assert
         expect(result).toEqual([]);
@@ -36,7 +36,7 @@ describe('queryContentModelBlocksBlocks', () => {
         };
 
         // Act
-        const result = queryContentModelBlocks(group, 'Table');
+        const result = queryContentModelBlocks<ReadonlyContentModelTable>(group, 'Table');
 
         // Assert
         expect(result).toEqual([]);