Merge pull request #115 from tushushu/wip-where-method

implement where method.

Author: tushushu

tests/test_base.py

@@ -286,6 +286,15 @@ def test_methods_no_arg(
         ('var', 'float', [1.0, 2.0, 3.0], 1.0, {"ddof": 1}),
         ('var', 'int', [1, 2, 3, 4], 1.25, {}),
         ('var', 'int', [1, 2, 3], 1.0, {"ddof": 1}),
+
+        ("where", "bool", [True, True, False, False], [
+         False, False], {"fn": lambda x: x == False},),  # noqa: E712
+        ("where", "float", [1.0, 2.0, 3.0, 4.0], [
+         1.0, 2.0], {"fn": lambda x: x < 3.0},),
+        ("where", "int", [1, 2, 3, 4], [
+            3, 4], {"fn": lambda x: x > 2},),
+        ("where", "string", ['foo', 'bar', 'baz'], [
+            'foo', 'baz'], {"fn": lambda x: x != 'bar'},),
     ],
 )
 def test_methods_with_args(

ulist/python/ulist/core.py

@@ -422,3 +422,30 @@ def var(self, ddof: int = 0) -> float:
         numerator = data.sub_scala(mean).pow_scala(2).sum()
         denominator = data.size() - ddof
         return numerator / denominator
+
+    def where(
+        self,
+        fn: Callable[[UltraFastList], UltraFastList]
+    ) -> UltraFastList:
+        """According to the function, return a ulist with elements of self
+        correspondingly.
+
+        Args:
+            fn (Callable[[UltraFastList], UltraFastList]):
+                Function to process self and return a BooleanList.
+
+        Returns:
+            UltraFastList: A ulist object.
+
+        Examples
+        --------
+        >>> import ulist as ul
+        >>> arr = ul.arange(6)
+        >>> arr
+        UltraFastList([0, 1, 2, 3, 4, 5])
+
+        >>> result = arr.where(lambda x: x > 2)
+        >>> result
+        UltraFastList([3, 4, 5])
+        """
+        return self.filter(fn(self))