Merge pull request #77 from tushushu/wip-select

Wip select

Author: tushushu

tests/test_constructor.py renamed to tests/test_constructors.py

@@ -0,0 +1,96 @@
+from typing import List, Union
+
+import pytest
+import ulist as ul
+import operator as op
+from ulist.utils import check_test_result
+
+LIST_TYPE = Union[List[float], List[int], List[bool]]
+
+
+@pytest.mark.parametrize(
+    "dtype, nums, kwargs, expected_value",
+    [
+        (
+            "float",
+            [0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
+            {
+                "conditions": [(op.lt, 2.0), (op.lt, 4.0)],
+                "choices": [False, True],
+                "default": False,
+            },
+            [False, False, True, True, False, False],
+        ),
+        (
+            "float",
+            [0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
+            {
+                "conditions": [(op.lt, 2.0), (op.lt, 4.0)],
+                "choices": [0.0, 1.0],
+                "default": 2.0,
+            },
+            [0.0, 0.0, 1.0, 1.0, 2.0, 2.0],
+        ),
+        (
+            "float",
+            [0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
+            {
+                "conditions": [(op.lt, 2.0), (op.lt, 4.0)],
+                "choices": [0, 1],
+                "default": 2,
+            },
+            [0, 0, 1, 1, 2, 2],
+        ),
+
+        (
+            "int",
+            [0, 1, 2, 3, 4, 5],
+            {
+                "conditions": [(op.lt, 2), (op.lt, 4)],
+                "choices": [False, True],
+                "default": False,
+            },
+            [False, False, True, True, False, False],
+        ),
+        (
+            "int",
+            [0, 1, 2, 3, 4, 5],
+            {
+                "conditions": [(op.lt, 2), (op.lt, 4)],
+                "choices": [0.0, 1.0],
+                "default": 2.0,
+            },
+            [0.0, 0.0, 1.0, 1.0, 2.0, 2.0],
+        ),
+        (
+            "int",
+            [0, 1, 2, 3, 4, 5],
+            {
+                "conditions": [(op.lt, 2), (op.lt, 4)],
+                "choices": [0, 1],
+                "default": 2,
+            },
+            [0, 0, 1, 1, 2, 2],
+        ),
+    ],
+)
+def test_select(
+    dtype: str,
+    nums: LIST_TYPE,
+    kwargs: dict,
+    expected_value: List[bool],
+) -> None:
+    arr = ul.from_seq(nums, dtype)
+    choices = kwargs["choices"]
+    default = kwargs["default"]
+    conditions = [f(arr, v) for f, v in kwargs["conditions"]]
+    result = ul.select(conditions, choices=choices, default=default)
+    if type(expected_value[0]) == int:
+        dtype = "int"
+    elif type(expected_value[0]) == float:
+        dtype = "float"
+    elif type(expected_value[0]) == bool:
+        dtype = "bool"
+    else:
+        raise TypeError(f"Unexpected type {type(expected_value[0])}!")
+    check_test_result(dtype, "ul.select", result, expected_value)

tests/test_control_flow.py

@@ -1,2 +1,3 @@
+from .control_flow import select  # noqa:F401
 from .constructor import arange, cycle, from_seq, repeat  # noqa:F401
 from .core import UltraFastList  # noqa:F401

ulist/python/ulist/__init__.py

@@ -0,0 +1,72 @@
+from typing import Callable, List
+
+from .core import BooleanList, UltraFastList
+from .typedef import ELEM, LIST_PY
+from .ulist import select_bool as _select_bool
+from .ulist import select_float as _select_float
+from .ulist import select_int as _select_int
+
+
+def select(
+        conditions: List[UltraFastList],
+        choices: LIST_PY,
+        default: ELEM,
+) -> UltraFastList:
+    """Return a ulist drawn from elements in `choices`, depending on`conditions`.
+
+    Args:
+        conditions (List[UltraFastList]):
+            The list of conditions which determine from which array in
+            `choices` the output elements are taken. When multiple conditions
+            are satisfied, the first one encountered in `conditions` is used.
+        choices (LIST_PY):
+            The list of ulist from which the output elements are taken.
+            It has to be of the same length as `conditions`.
+        default (ELEM):
+            The element inserted in output when all conditions evaluate
+            to False.
+
+    Raises:
+        TypeError:
+            The type of parameter `default` should be bool, float or int!
+
+    Returns:
+        UltraFastList: A ulist object.
+
+    Examples
+    --------
+    >>> import ulist as ul
+    >>> arr = ul.arange(6)
+    >>> arr
+    UltraFastList([0, 1, 2, 3, 4, 5])
+
+    >>> conditions = [arr < 2, arr < 4]
+    >>> conditions
+    [
+        UltraFastList([True, True, False, False, False, False]),
+        UltraFastList([True, True, True, True, False, False])
+    ]
+
+    >>> result = ul.select(conditions, choices=[0, 1], default=2)
+    >>> result
+    UltraFastList([0, 0, 1, 1, 2, 2])
+    """
+    assert len(conditions) == len(choices)
+
+    if type(default) is bool:
+        fn: Callable = _select_bool
+    elif type(default) is float:
+        fn = _select_float
+    elif type(default) is int:
+        fn = _select_int
+    else:
+        raise TypeError(
+            "The type of parameter `default` should be bool, float or int!"
+        )
+
+    _conditions = []
+    for cond in conditions:
+        assert isinstance(cond._values, BooleanList)
+        _conditions.append(cond._values)
+    result = fn(_conditions, choices, default)
+    return UltraFastList(result)

ulist/python/ulist/control_flow.py

@@ -1,6 +1,6 @@
-from typing import Sequence, List
+from typing import List, Sequence
 
-from .typedef import ELEM, NUM, NUM_LIST_RS
+from .typedef import ELEM, LIST_PY, NUM, NUM_LIST_RS
 
 
 class BooleanList:
@@ -115,3 +115,24 @@ class IntegerList:
 
 
 def arange(start: int, stop: int, step: int) -> IntegerList: ...
+
+
+def select_bool(
+    conditions: List[BooleanList],
+    choices: LIST_PY,
+    default: bool,
+) -> BooleanList: ...
+
+
+def select_float(
+    conditions: List[BooleanList],
+    choices: LIST_PY,
+    default: float,
+) -> FloatList: ...
+
+
+def select_int(
+    conditions: List[BooleanList],
+    choices: LIST_PY,
+    default: int,
+) -> IntegerList: ...

ulist/python/ulist/ulist.pyi

@@ -0,0 +1,61 @@
+use crate::base::List;
+use crate::boolean::BooleanList;
+use crate::float::FloatList;
+use crate::integer::IntegerList;
+use pyo3::prelude::*;
+use pyo3::Py;
+use std::cmp::PartialEq;
+use std::marker::Copy;
+
+unsafe fn select<T, U>(
+    py: Python,
+    conditions: &Vec<Py<BooleanList>>,
+    choices: &Vec<T>,
+    default: T,
+) -> U
+where
+    T: PartialEq + Copy,
+    U: List<T>,
+{
+    let cond: Vec<PyRef<BooleanList>> = conditions.iter().map(|x| x.borrow(py)).collect();
+    let mut vec = vec![default; cond[0].size()];
+    for j in 0..cond[0].size() {
+        for i in 0..cond.len() {
+            if cond[i].get(j) {
+                vec[j] = choices[i];
+                break;
+            }
+        }
+    }
+    U::_new(vec)
+}
+
+#[pyfunction]
+pub unsafe fn select_bool(
+    py: Python,
+    conditions: Vec<Py<BooleanList>>,
+    choices: Vec<bool>,
+    default: bool,
+) -> BooleanList {
+    select::<bool, BooleanList>(py, &conditions, &choices, default)
+}
+
+#[pyfunction]
+pub unsafe fn select_float(
+    py: Python,
+    conditions: Vec<Py<BooleanList>>,
+    choices: Vec<f32>,
+    default: f32,
+) -> FloatList {
+    select::<f32, FloatList>(py, &conditions, &choices, default)
+}
+
+#[pyfunction]
+pub unsafe fn select_int(
+    py: Python,
+    conditions: Vec<Py<BooleanList>>,
+    choices: Vec<i32>,
+    default: i32,
+) -> IntegerList {
+    select::<i32, IntegerList>(py, &conditions, &choices, default)
+}

ulist/src/control_flow.rs

@@ -1,17 +1,18 @@
 mod base;
 mod boolean;
+mod control_flow;
 mod float;
 mod integer;
 mod numerical;
 mod types;
-use integer::IntegerList;
+use control_flow::*;
 use pyo3::prelude::*;
 
 // Cannot find a way to put this function in another file.
 #[pyfunction]
-pub fn arange(start: i32, stop: i32, step: usize) -> IntegerList {
+pub fn arange(start: i32, stop: i32, step: usize) -> integer::IntegerList {
     let vec = (start..stop).step_by(step).collect();
-    IntegerList::new(vec)
+    integer::IntegerList::new(vec)
 }
 
 /// A Python module implemented in Rust. The name of this function must match
@@ -23,6 +24,9 @@ fn ulist(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<float::FloatList>()?;
     m.add_class::<integer::IntegerList>()?;
     m.add_function(wrap_pyfunction!(arange, m)?)?;
+    m.add_function(wrap_pyfunction!(select_bool, m)?)?;
+    m.add_function(wrap_pyfunction!(select_float, m)?)?;
+    m.add_function(wrap_pyfunction!(select_int, m)?)?;
 
     Ok(())
 }