Skip to content

Commit

Permalink
further test standartisation
Browse files Browse the repository at this point in the history
  • Loading branch information
@ohrechykha
ohrechykha committed Aug 28, 2024
1 parent f04ccbf commit 7b30c58
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 40 deletions.
28 changes: 15 additions & 13 deletions src/ragged/_spec_set_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
if isinstance(x, ragged.array):
if x.ndim == 0:
return unique_all_result(
values=ragged.array([x]),
values=ragged.array(x),
indices=ragged.array([0]),
inverse_indices=ragged.array([0]),
counts=ragged.array([1]),
Expand All @@ -76,8 +76,8 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
counts=ragged.array(counts),
)
else:
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore
msg = f"Expected ragged type but got {type(x)}" # type: ignore[unreachable]
raise TypeError(msg)


unique_counts_result = namedtuple( # pylint: disable=C0103
Expand Down Expand Up @@ -109,7 +109,7 @@ def unique_counts(x: array, /) -> tuple[array, array]:
if isinstance(x, ragged.array):
if x.ndim == 0:
return unique_counts_result(
values=ragged.array([x]), counts=ragged.array([1])
values=ragged.array(x), counts=ragged.array([1])
)
else:
x_flat = ak.ravel(x._impl)
Expand All @@ -118,8 +118,8 @@ def unique_counts(x: array, /) -> tuple[array, array]:
values=ragged.array(values), counts=ragged.array(counts)
)
else:
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore
msg = f"Expected ragged type but got {type(x)}" # type: ignore[unreachable]
raise TypeError(msg)


unique_inverse_result = namedtuple( # pylint: disable=C0103
Expand Down Expand Up @@ -149,8 +149,10 @@ def unique_inverse(x: array, /) -> tuple[array, array]:
https://data-apis.org/array-api/latest/API_specification/generated/array_api.unique_inverse.html
"""
if isinstance(x, ragged.array):
if ak.is_scalar(x):
return unique_inverse_result(values=x, inverse_indices=ragged.array([0]))
if x.ndim == 0:
return unique_inverse_result(
values=ragged.array(x), inverse_indices=ragged.array([0])
)
else:
x_flat = ak.ravel(x._impl)
values, inverse_indices = np.unique(x_flat.layout.data, return_inverse=True)
Expand All @@ -160,8 +162,8 @@ def unique_inverse(x: array, /) -> tuple[array, array]:
inverse_indices=ragged.array(inverse_indices),
)
else:
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore
msg = f"Expected ragged type but got {type(x)}" # type: ignore[unreachable]
raise TypeError(msg)


def unique_values(x: array, /) -> array:
Expand All @@ -181,11 +183,11 @@ def unique_values(x: array, /) -> array:
"""
if isinstance(x, ragged.array):
if x.ndim == 0:
return ragged.array([x])
return ragged.array(x)

else:
x_flat = ak.ravel(x._impl)
return ragged.array(np.unique(x_flat.layout.data))
else:
err = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(err) # type: ignore
err = f"Expected ragged type but got {type(x)}" # type: ignore[unreachable]
raise TypeError(err)
56 changes: 29 additions & 27 deletions tests/test_spec_set_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from __future__ import annotations

import awkward as ak
import pytest

import ragged

Expand All @@ -26,22 +25,25 @@ def test_existence():


def test_can_take_list():
assert ak.to_list(
ragged.unique_values(ragged.array([1, 2, 4, 3, 4, 5, 6, 20]))
) == ak.to_list(ragged.array([1, 2, 3, 4, 5, 6, 20]))
arr = ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
expected_unique_values = ragged.array([1, 2, 3, 4, 5, 6, 20])
unique_values = ragged.unique_values(arr)
assert ak.to_list(expected_unique_values) == ak.to_list(unique_values)


def test_can_take_empty_arr():
# with pytest.raises(TypeError):
assert ak.to_list(ragged.unique_values(ragged.array([]))) == ak.to_list(
ragged.array([])
)
arr = ragged.array([])
expected_unique_values = ragged.array([0])
unique_values = ragged.unique_values(arr)
assert ak.to_list(expected_unique_values) == ak.to_list(unique_values)


def test_can_take_moredimensions():
assert ak.to_list(
ragged.unique_values(ragged.array([[1, 2, 2, 3, 4], [5, 6]]))
) == ak.to_list(ragged.array([1, 2, 3, 4, 5, 6]))
arr = ragged.array([[1, 2, 2, 3, 4], [5, 6]])
expected_unique_values = ragged.array([1, 2, 3, 4, 5, 6])
unique_values = ragged.unique_values(arr)
assert ak.to_list(expected_unique_values) == ak.to_list(unique_values)


def test_can_take_1d_array():
Expand Down Expand Up @@ -135,23 +137,23 @@ def test_can_inverse_scalar():


# unique_all tests
def test_can_all_none():
with pytest.raises(TypeError):
arr = None
expected_unique_values = ragged.array(None)
expected_unique_indices = ragged.array(None)
expected_unique_inverse = ragged.array(None)
expected_unique_counts = ragged.array(None)
(
unique_values,
unique_indices,
unique_inverse,
unique_counts,
) = ragged.unique_all(arr)
assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
assert ak.to_list(unique_inverse) == ak.to_list(expected_unique_inverse)
assert ak.to_list(unique_counts) == ak.to_list(expected_unique_counts)
# def test_can_all_none():
# with pytest.raises(TypeError):
# arr = None
# expected_unique_values = ragged.array(None)
# expected_unique_indices = ragged.array(None)
# expected_unique_inverse = ragged.array(None)
# expected_unique_counts = ragged.array(None)
# (
# unique_values,
# unique_indices,
# unique_inverse,
# unique_counts,
# ) = ragged.unique_all(arr)
# assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
# assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
# assert ak.to_list(unique_inverse) == ak.to_list(expected_unique_inverse)
# assert ak.to_list(unique_counts) == ak.to_list(expected_unique_counts)


def test_can_all_list():
Expand Down

0 comments on commit 7b30c58

Please sign in to comment.