Fix types

Author: satya164

src/__tests__/rehype-static-to-dynamic.test.ts

@@ -1,22 +1,30 @@
 import dedent from 'dedent';
 import assert from 'node:assert';
 import { describe, test } from 'node:test';
-import rehypeStaticToDynamic from '../plugins/rehype-static-to-dynamic.ts';
+import rehypeStaticToDynamic, {
+  type RehypeStaticToDynamicElement as Element,
+  type RehypeStaticToDynamicElementChild as ElementChild,
+  type RehypeStaticToDynamicRoot as Root,
+  type RehypeStaticToDynamicText as Text,
+  type RehypeStaticToDynamicTreeChild as TreeChild,
+} from '../plugins/rehype-static-to-dynamic.ts';
 
 /**
  * Helper function to create a test tree structure
  */
-function createTestTree(code) {
-  return {
+function createTestTree(code: string): Root {
+  const tree: Root = {
     type: 'root',
     children: [
       {
         type: 'element',
         tagName: 'pre',
+        properties: {},
         children: [
           {
             type: 'element',
             tagName: 'code',
+            properties: {},
             data: { meta: 'static2dynamic' },
             children: [
               {
@@ -29,23 +37,25 @@ function createTestTree(code) {
       },
     ],
   };
+
+  return tree;
 }
 
 /**
  * Helper function to extract the transformed code from the tree
  */
-function extractTransformedCode(tree) {
+function extractTransformedCode(tree: Root): string {
   // After transformation, the tree should have TabItem elements
   const tabsElement = tree.children[0];
 
-  if (!tabsElement || tabsElement.tagName !== 'Tabs') {
+  if (!isElement(tabsElement) || tabsElement.tagName !== 'Tabs') {
     throw new Error('Expected Tabs element not found');
   }
 
   // Find the "Dynamic" tab
   const dynamicTab = tabsElement.children.find(
-    (child) =>
-      child.type === 'element' &&
+    (child): child is Element =>
+      isElement(child) &&
       child.tagName === 'TabItem' &&
       child.properties?.value === 'dynamic'
   );
@@ -56,26 +66,38 @@ function extractTransformedCode(tree) {
 
   // Extract the code from the dynamic tab
   const preElement = dynamicTab.children.find(
-    (child) => child.type === 'element' && child.tagName === 'pre'
+    (child): child is Element =>
+      isElement(child) && child.tagName === 'pre'
   );
 
   if (!preElement) {
     throw new Error('Pre element not found in dynamic tab');
   }
 
   const codeElement = preElement.children.find(
-    (child) => child.type === 'element' && child.tagName === 'code'
+    (child): child is Element =>
+      isElement(child) && child.tagName === 'code'
   );
 
   if (!codeElement) {
     throw new Error('Code element not found');
   }
 
-  const textNode = codeElement.children.find((child) => child.type === 'text');
+  const textNode = codeElement.children.find(isTextNode);
 
   return textNode?.value || '';
 }
 
+function isElement(
+  node: TreeChild | ElementChild | undefined
+): node is Element {
+  return Boolean(node && node.type === 'element');
+}
+
+function isTextNode(node: ElementChild): node is Text {
+  return node.type === 'text';
+}
+
 describe('rehype-static-to-dynamic', () => {
   test('basic screen transformation', async () => {
     const input = dedent /* javascript */ `

src/plugins/rehype-static-to-dynamic.ts

@@ -8,7 +8,7 @@ import * as recast from 'recast';
 import * as babelParser from 'recast/parsers/babel-ts.js';
 import { visit } from 'unist-util-visit';
 import { fileURLToPath } from 'url';
-import type { Element, Root } from 'hast';
+import type { Element, Root, Text } from 'hast';
 import type { Parent } from 'unist';
 
 const __filename = fileURLToPath(import.meta.url);
@@ -22,6 +22,16 @@ type TrackedComment = {
   type: t.Comment['type'];
 };
 
+export type RehypeStaticToDynamicRoot = Root;
+
+export type RehypeStaticToDynamicElement = Element;
+
+export type RehypeStaticToDynamicText = Text;
+
+export type RehypeStaticToDynamicTreeChild = Root['children'][number];
+
+export type RehypeStaticToDynamicElementChild = Element['children'][number];
+
 type CommentWithMarkers = t.Comment & {
   leading?: boolean;
   trailing?: boolean;