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 committed Dec 12, 2024
1 parent c2c3b03 commit 3060b64
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 3060b64

Please sign in to comment.