Types/improve operator node types (josdejong#2576)

mattvague · web-flow · commit cd33da8d16aa · 2022-05-31T11:41:32.000+02:00
* Add generic args to OperatorNode for `fn` and `op`

* Fix lint

* Format and add generic for args as well

* Make OperatorNode generics optional using `never`
diff --git a/types/index.d.ts b/types/index.d.ts
@@ -276,25 +276,67 @@ declare namespace math {
     new (properties: Record<string, MathNode>): ObjectNode
   }
 
-  interface OperatorNode extends MathNodeCommon {
+  type OperatorNodeMap = {
+    xor: 'xor'
+    and: 'and'
+    bitOr: '|'
+    bitXor: '^|'
+    bitAnd: '&'
+    equal: '=='
+    unequal: '!='
+    smaller: '<'
+    larger: '>'
+    smallerEq: '<='
+    leftShift: '<<'
+    rightArithShift: '>>'
+    rightLogShift: '>>>'
+    to: 'to'
+    add: '+'
+    subtract: '-'
+    multiply: '*'
+    divide: '/'
+    dotMultiply: '.*'
+    dotDivide: './'
+    mod: 'mod'
+    unaryPlus: '+'
+    unaryMinus: '-'
+    bitNot: '~'
+    not: 'not'
+    pow: '^'
+    dotPow: '.^'
+    factorial: '!'
+  }
+
+  type OperatorNodeOp = OperatorNodeMap[keyof OperatorNodeMap]
+  type OperatorNodeFn = keyof OperatorNodeMap
+
+  interface OperatorNode<
+    TOp extends OperatorNodeMap[TFn] = never,
+    TFn extends OperatorNodeFn = never,
+    TArgs extends MathNode[] = MathNode[]
+  > extends MathNodeCommon {
     type: 'OperatorNode'
     isOperatorNode: true
-    op: string
-    fn: string
-    args: MathNode[]
+    op: TOp
+    fn: TFn
+    args: TArgs
     implicit: boolean
     isUnary(): boolean
     isBinary(): boolean
   }
-  interface OperatorNodeCtor {
-    new (
-      op: string,
-      fn: string,
-      args: MathNode[],
+
+  interface OperatorNodeCtor extends MathNodeCommon {
+    new <
+      TOp extends OperatorNodeMap[TFn],
+      TFn extends OperatorNodeFn,
+      TArgs extends MathNode[]
+    >(
+      op: TOp,
+      fn: TFn,
+      args: TArgs,
       implicit?: boolean
-    ): OperatorNode
+    ): OperatorNode<TOp, TFn, TArgs>
   }
-
   interface ParenthesisNode extends MathNodeCommon {
     type: 'ParenthesisNode'
     isParenthesisNode: true
@@ -345,7 +387,7 @@ declare namespace math {
     | FunctionNode
     | IndexNode
     | ObjectNode
-    | OperatorNode
+    | OperatorNode<OperatorNodeOp, OperatorNodeFn>
     | ParenthesisNode
     | RangeNode
     | RelationalNode
@@ -3095,7 +3137,9 @@ declare namespace math {
 
     isObjectNode(x: unknown): x is ObjectNode
 
-    isOperatorNode(x: unknown): x is OperatorNode
+    isOperatorNode(
+      x: unknown
+    ): x is OperatorNode<OperatorNodeOp, OperatorNodeFn>
 
     isParenthesisNode(x: unknown): x is ParenthesisNode
 
diff --git a/types/index.ts b/types/index.ts
@@ -25,6 +25,9 @@ import {
   MathNumericType,
   ConstantNode,
   OperatorNode,
+  OperatorNodeFn,
+  OperatorNodeOp,
+  SymbolNode,
 } from 'mathjs'
 import * as assert from 'assert'
 import { expectTypeOf } from 'expect-type'
@@ -1527,23 +1530,34 @@ Function round examples
   */
 {
   const math = create(all, {})
+  expectTypeOf(
+    new math.OperatorNode('/', 'divide', [
+      new math.ConstantNode(3),
+      new math.SymbolNode('x'),
+    ])
+  ).toMatchTypeOf<OperatorNode<'/', 'divide', (ConstantNode | SymbolNode)[]>>()
+
   expectTypeOf(new math.ConstantNode(1).clone()).toMatchTypeOf<ConstantNode>()
   expectTypeOf(
     new math.OperatorNode('*', 'multiply', [
       new math.ConstantNode(3),
       new math.SymbolNode('x'),
     ]).clone()
-  ).toMatchTypeOf<OperatorNode>()
+  ).toMatchTypeOf<
+    OperatorNode<'*', 'multiply', (ConstantNode | SymbolNode)[]>
+  >()
 
   expectTypeOf(
     new math.ConstantNode(1).cloneDeep()
   ).toMatchTypeOf<ConstantNode>()
   expectTypeOf(
-    new math.OperatorNode('*', 'multiply', [
+    new math.OperatorNode('+', 'unaryPlus', [
       new math.ConstantNode(3),
       new math.SymbolNode('x'),
     ]).cloneDeep()
-  ).toMatchTypeOf<OperatorNode>()
+  ).toMatchTypeOf<
+    OperatorNode<'+', 'unaryPlus', (ConstantNode | SymbolNode)[]>
+  >()
 
   expectTypeOf(
     math.clone(new math.ConstantNode(1))
@@ -1900,7 +1914,9 @@ Factory Test
     expectTypeOf(x).toMatchTypeOf<math.ObjectNode>()
   }
   if (math.isOperatorNode(x)) {
-    expectTypeOf(x).toMatchTypeOf<math.OperatorNode>()
+    expectTypeOf(x).toMatchTypeOf<
+      OperatorNode<OperatorNodeOp, OperatorNodeFn, MathNode[]>
+    >()
   }
   if (math.isParenthesisNode(x)) {
     expectTypeOf(x).toMatchTypeOf<math.ParenthesisNode>()
