From 25573bfa66bedc38b60dc39575ea7d17a495edcc Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Thu, 12 Dec 2024 09:43:31 -0800
Subject: [PATCH] Fixes numpy annotation test

So new version of numpy doesn't have the same behavior for 3.9 as
for 3.10+.
---
 tests/test_node.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_node.py b/tests/test_node.py
index ed8c96aa2..69ba32990 100644
--- a/tests/test_node.py
+++ b/tests/test_node.py
@@ -60,7 +60,7 @@ def annotated_func(first: ArrayN[np.float64], other: float = 2.0) -> ArrayN[np.f
 
     node = Node.from_fn(annotated_func)
     assert node.name == "annotated_func"
-    if major == 2 and minor > 1:  # greater that 2.1
+    if major == 2 and minor > 1 and sys.version_info > (3, 9):  # greater that 2.1
         expected = {
             "first": (
                 Annotated[np.ndarray[tuple[int, ...], np.dtype[np.float64]], Literal["N"]],
@@ -68,6 +68,7 @@ def annotated_func(first: ArrayN[np.float64], other: float = 2.0) -> ArrayN[np.f
             ),
             "other": (float, DependencyType.OPTIONAL),
         }
+        expected_type = Annotated[np.ndarray[tuple[int, ...], np.dtype[np.float64]], Literal["N"]]
     else:
         expected = {
             "first": (
@@ -76,8 +77,9 @@ def annotated_func(first: ArrayN[np.float64], other: float = 2.0) -> ArrayN[np.f
             ),
             "other": (float, DependencyType.OPTIONAL),
         }
+        expected_type = Annotated[np.ndarray[Any, np.dtype[np.float64]], Literal["N"]]
     assert node.input_types == expected
-    assert node.type == Annotated[np.ndarray[tuple[int, ...], np.dtype[np.float64]], Literal["N"]]
+    assert node.type == expected_type
 
 
 @pytest.mark.parametrize(