Skip to content

Commit

Permalink
Fixes numpy annotation test
Browse files Browse the repository at this point in the history
So new version of numpy doesn't have the same behavior for 3.9 as
for 3.10+.
  • Loading branch information
skrawcz authored and elijahbenizzy committed Dec 12, 2024
1 parent 05ec8bc commit 25573bf
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions tests/test_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,14 +60,15 @@ 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"]],
DependencyType.REQUIRED,
),
"other": (float, DependencyType.OPTIONAL),
}
expected_type = Annotated[np.ndarray[tuple[int, ...], np.dtype[np.float64]], Literal["N"]]
else:
expected = {
"first": (
Expand All @@ -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(
Expand Down

0 comments on commit 25573bf

Please sign in to comment.