Add nested type guards

metaplex-foundation · Apr 9, 2024 · dbc6281 · dbc6281
1 parent 1eb39c8
commit dbc6281
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 0 deletions.
diff --git a/src/nodes/typeNodes/NestedTypeNode.ts b/src/nodes/typeNodes/NestedTypeNode.ts
@@ -54,3 +54,15 @@ export function isNestedTypeNode<TKind extends TypeNode['kind']>(
   const resolved = resolveNestedTypeNode(node);
   return !!node && kinds.includes(resolved.kind as TKind);
 }
+
+export function assertIsNestedTypeNode<TKind extends TypeNode['kind']>(
+  node: Node | null | undefined,
+  kind: TKind | TKind[]
+): asserts node is NestedTypeNode<Extract<TypeNode, { kind: TKind }>> {
+  const kinds = Array.isArray(kind) ? kind : [kind];
+  if (!isNestedTypeNode(node, kinds)) {
+    throw new Error(
+      `Expected nested type of ${kinds.join(' | ')}, got ${node?.kind ?? 'null'}.`
+    );
+  }
+}
diff --git a/test/nodes/typeNodes/NestedTypeNode.test.ts b/test/nodes/typeNodes/NestedTypeNode.test.ts
@@ -1,5 +1,6 @@
 import test from 'ava';
 import {
+  assertIsNestedTypeNode,
   fixedSizeTypeNode,
   isNestedTypeNode,
   numberTypeNode,
@@ -52,3 +53,16 @@ test('it checks if a node is a nested type', (t) => {
   t.true(isNestedTypeNode(nestedNode, 'numberTypeNode'));
   t.false(isNestedTypeNode(nestedNode, 'stringTypeNode'));
 });
+
+test('it asserts that a node is a nested type', (t) => {
+  const flatNode = numberTypeNode('u64');
+  t.notThrows(() => assertIsNestedTypeNode(flatNode, 'numberTypeNode'));
+  t.throws(() => assertIsNestedTypeNode(flatNode, 'stringTypeNode'));
+
+  const nestedNode = sizePrefixTypeNode(
+    fixedSizeTypeNode(numberTypeNode('u64'), 32),
+    numberTypeNode('u8')
+  );
+  t.notThrows(() => assertIsNestedTypeNode(nestedNode, 'numberTypeNode'));
+  t.throws(() => assertIsNestedTypeNode(nestedNode, 'stringTypeNode'));
+});
diff --git a/test/nodes/typeNodes/NestedTypeNode.typetest.ts b/test/nodes/typeNodes/NestedTypeNode.typetest.ts
@@ -6,6 +6,10 @@ import {
   NestedTypeNode,
   StringTypeNode,
   stringTypeNode,
+  transformNestedTypeNode,
+  Node,
+  isNestedTypeNode,
+  assertIsNestedTypeNode,
 } from '../../../src';
 
 {
@@ -31,3 +35,47 @@ import {
   // @ts-expect-error The nested type is not a string.
   resolveNestedTypeNode(numberNestedNode) satisfies StringTypeNode;
 }
+
+{
+  // [transformNestedTypeNode]: it transforms the nested type of a nested node.
+  const transformedNode = transformNestedTypeNode(
+    fixedSizeTypeNode(stringTypeNode(), 32),
+    () => numberTypeNode('u32')
+  );
+  transformedNode satisfies NestedTypeNode<NumberTypeNode>;
+  // @ts-expect-error The nested type is not a number.
+  transformedNode satisfies NestedTypeNode<StringTypeNode>;
+}
+
+{
+  // [isNestedTypeNode]: it narrows the type of a node to a nested type node.
+  const node = {} as Node;
+  if (isNestedTypeNode(node, 'numberTypeNode')) {
+    node satisfies NestedTypeNode<NumberTypeNode>;
+    // @ts-expect-error The nested type is not a string.
+    node satisfies NestedTypeNode<StringTypeNode>;
+  }
+  if (isNestedTypeNode(node, 'stringTypeNode')) {
+    node satisfies NestedTypeNode<StringTypeNode>;
+    // @ts-expect-error The nested type is not a number.
+    node satisfies NestedTypeNode<NumberTypeNode>;
+  }
+}
+
+{
+  // [assertIsNestedTypeNode]: it narrows the type of a node to a nested type node.
+  {
+    const node = {} as Node;
+    assertIsNestedTypeNode(node, 'numberTypeNode');
+    node satisfies NestedTypeNode<NumberTypeNode>;
+    // @ts-expect-error The nested type is not a string.
+    node satisfies NestedTypeNode<StringTypeNode>;
+  }
+  {
+    const node = {} as Node;
+    assertIsNestedTypeNode(node, 'stringTypeNode');
+    node satisfies NestedTypeNode<StringTypeNode>;
+    // @ts-expect-error The nested type is not a number.
+    node satisfies NestedTypeNode<NumberTypeNode>;
+  }
+}