From 6ec236e0b70e47b9f8e03a3f0ac54d872c3c0fb4 Mon Sep 17 00:00:00 2001
From: Justin DuJardin <justin@dujardinconsulting.com>
Date: Fri, 15 Dec 2023 12:17:32 -0800
Subject: [PATCH] fix(typing): proper generic type for BinaryTreeNode

 - pyright finally came calling on our swapping of left/right/parent fields in the MathExpression superclass.
---
 mathy_core/expressions.py |  8 ++++----
 mathy_core/tree.py        | 28 ++++++++++++++--------------
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/mathy_core/expressions.py b/mathy_core/expressions.py
index 624b222..d831d66 100644
--- a/mathy_core/expressions.py
+++ b/mathy_core/expressions.py
@@ -28,7 +28,7 @@
 MathTypeKeysMax = max(MathTypeKeys.values()) + 1
 
 
-class MathExpression(BinaryTreeNode):
+class MathExpression(BinaryTreeNode["MathExpression"]):
     """Math tree node with helpers for manipulating expressions.
 
     `mathy:x+y=z`
@@ -117,7 +117,7 @@ def visit_fn(
     def with_color(self, text: str, style: str = "bright") -> str:
         """Render a string that is colored if something has changed"""
         if self._rendering_change is True and self._changed is True:
-            return color(text, fore=self.color, style=style)
+            return f"{color(text, fore=self.color, style=style)}"
         return text
 
     def add_class(self, classes: Union[List[str], str]) -> "MathExpression":
@@ -147,7 +147,7 @@ def visit_fn(
 
     def to_list(self, visit: str = "preorder") -> List["MathExpression"]:
         """Convert this node hierarchy into a list."""
-        results = []
+        results: List[MathExpression] = []
 
         def visit_fn(
             node: MathExpression, depth: int, data: Any
@@ -688,7 +688,7 @@ def clone(self) -> "ConstantExpression":  # type:ignore[override]
         result.value = self.value
         return result  # type:ignore
 
-    def evaluate(self, _context: Optional[Dict[str, NumberType]] = None) -> NumberType:
+    def evaluate(self, context: Optional[Dict[str, NumberType]] = None) -> NumberType:
         assert self.value is not None
         return self.value
 
diff --git a/mathy_core/tree.py b/mathy_core/tree.py
index d91070c..51f4d67 100644
--- a/mathy_core/tree.py
+++ b/mathy_core/tree.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, List, Optional, TypeVar, Union, cast
+from typing import Any, Callable, Generic, List, Optional, TypeVar, Union, cast
 
 from .types import Literal
 
@@ -29,7 +29,7 @@
 ]
 
 
-class BinaryTreeNode:
+class BinaryTreeNode(Generic[NodeType]):
     """
     The binary tree node is the base node for all of our trees, and provides a
     rich set of methods for constructing, inspecting, and modifying them.
@@ -44,18 +44,18 @@ class BinaryTreeNode:
     y: Optional[float]
     offset: Optional[float]
     level: Optional[int]
-    thread: Optional["BinaryTreeNode"]
+    thread: Optional[NodeType]
 
-    left: Optional["BinaryTreeNode"]
-    right: Optional["BinaryTreeNode"]
-    parent: Optional["BinaryTreeNode"]
+    left: Optional[NodeType]
+    right: Optional[NodeType]
+    parent: Optional[NodeType]
 
     #  Allow specifying children in the constructor
     def __init__(
-        self,
-        left: Optional["BinaryTreeNode"] = None,
-        right: Optional["BinaryTreeNode"] = None,
-        parent: Optional["BinaryTreeNode"] = None,
+        self: NodeType,
+        left: Optional[NodeType] = None,
+        right: Optional[NodeType] = None,
+        parent: Optional[NodeType] = None,
         id: Optional[str] = None,
     ):
         if id is None:
@@ -208,7 +208,7 @@ def get_root(self: NodeType) -> NodeType:
 
         return cast(NodeType, result)
 
-    def get_root_side(self: "BinaryTreeNode") -> Literal["left", "right"]:
+    def get_root_side(self: NodeType) -> Literal["left", "right"]:
         """Return the side of the tree that this node lives on"""
         result = self
         last_child = None
@@ -225,7 +225,7 @@ def get_root_side(self: "BinaryTreeNode") -> Literal["left", "right"]:
 
     def set_left(
         self: NodeType,
-        child: Optional["BinaryTreeNode"] = None,
+        child: Optional[NodeType] = None,
         clear_old_child_parent: bool = False,
     ) -> NodeType:
         """Set the left node to the passed `child`"""
@@ -241,7 +241,7 @@ def set_left(
 
     def set_right(
         self: NodeType,
-        child: Optional["BinaryTreeNode"] = None,
+        child: Optional[NodeType] = None,
         clear_old_child_parent: bool = False,
     ) -> NodeType:
         """Set the right node to the passed `child`"""
@@ -255,7 +255,7 @@ def set_right(
 
         return self
 
-    def get_side(self, child: Optional["BinaryTreeNode"]) -> Literal["left", "right"]:
+    def get_side(self, child: Optional[NodeType]) -> Literal["left", "right"]:
         """Determine whether the given `child` is the left or right child of this
         node"""
         if child == self.left: