🗓 Jun 30, 2024 2:45:57 PM

✨ without method
✨ pick method
🧪 tests added/updated
🤖 types added/updated

chepy/modules/utils.py

@@ -2,7 +2,7 @@
 import random
 import difflib
 from collections import OrderedDict
-from typing import TypeVar, Union
+from typing import TypeVar, Union, Any
 
 import chepy.modules.internal.colors as _int_colors
 
@@ -252,7 +252,9 @@ def split_chunks(self, chunk_size) -> UtilsT:
         Returns:
             Chepy: The Chepy object.
         """
-        data = self._convert_to_bytes()
+        data = self.state
+        if not isinstance(data, list):
+            data = self._convert_to_bytes()
         data_chunks = []
         for i in range(0, len(data), chunk_size):
             data_chunks.append(data[i : i + chunk_size])
@@ -747,3 +749,73 @@ def drop_bytes(self, start: int, length: int) -> UtilsT:
         data = self._convert_to_bytes()
         self.state = data[:start] + data[end:]
         return self
+
+    @ChepyDecorators.call_stack
+    def without(self, *values: Any):
+        """Remove specified values from the state. Works on strings, bytes, lists and dicts
+
+        Raises:
+            TypeError: If state does not contain valid data
+
+        Returns:
+            Chepy: The chepy object.
+        """
+        collection = self.state
+        if isinstance(collection, list):
+            self.state = [item for item in collection if item not in values]
+        elif isinstance(collection, dict):
+            self.state = {
+                k: v
+                for k, v in collection.items()
+                if k not in values and v not in values
+            }
+        elif isinstance(
+            collection,
+            (
+                bytes,
+                str,
+            ),
+        ):
+            if isinstance(collection, str):
+                collection = collection.encode()
+            byte_values = set()
+            for value in values:
+                if isinstance(value, str):
+                    byte_values.update(value.encode())
+                elif isinstance(value, bytes):
+                    byte_values.update(value)
+            self.state = bytes([char for char in collection if char not in byte_values])
+        else:  # pragma: no cover
+            raise TypeError("Input should be a list, dictionary, string, or bytes")
+        return self
+
+    @ChepyDecorators.call_stack
+    def pick(self, *values: Any):
+        """Only pick specified values from the state. Works on strings, bytes, lists and dicts
+
+        Raises:
+            TypeError: If state does not contain valid data
+
+        Returns:
+            Chepy: The chepy object.
+        """
+        collection = self.state
+        if isinstance(collection, list):
+            self.state = [item for item in collection if item in values]
+        elif isinstance(collection, dict):
+            self.state = {
+                k: v for k, v in collection.items() if k in values or v in values
+            }
+        elif isinstance(collection, (bytes, str)):
+            if isinstance(collection, str):
+                collection = collection.encode()
+            byte_values = set()
+            for value in values:
+                if isinstance(value, str):
+                    byte_values.update(value.encode())
+                elif isinstance(value, bytes):
+                    byte_values.update(value)
+            self.state = bytes([char for char in collection if char in byte_values])
+        else:  # pragma: no cover
+            raise TypeError("Input should be a list, dictionary, string, or bytes")
+        return self

chepy/modules/utils.pyi

@@ -39,3 +39,5 @@ class Utils(ChepyCore):
     def regex_to_str(self: UtilsT, all_combo: bool=...) -> UtilsT: ...
     def shuffle(self: UtilsT) -> UtilsT: ...
     def drop_bytes(self: UtilsT, start: int, length: int) -> UtilsT: ...
+    def without(self: UtilsT, *values: Any) -> UtilsT: ...
+    def pick(self: UtilsT, *values: Any) -> UtilsT: ...

tests/test_utils.py

@@ -32,11 +32,7 @@ def test_search():
         == 5
     )
     assert (
-        len(
-            Chepy("loLolololoL")
-            .regex_search("ol", ignore_case=True, is_bytes=True)
-            .o
-        )
+        len(Chepy("loLolololoL").regex_search("ol", ignore_case=True, is_bytes=True).o)
         == 5
     )
     assert (
@@ -102,6 +98,7 @@ def test_split_chunks():
     c1 = Chepy(data).split_chunks(2)
     assert len(c1.o) == 6
     assert c1.o == [b"he", b"ll", b"o ", b"wo", b"rl", b"d"]
+    assert Chepy([1, 2, 3, 4]).split_chunks(2).o == [[1, 2], [3, 4]]
 
 
 def test_select_n():
@@ -248,3 +245,18 @@ def test_shuffle():
 
 def test_drop_bytes():
     assert Chepy("hello").drop_bytes(2, 2).o == b"heo"
+
+def test_without_pick():
+    data1 = 'hello'
+    data2 = [1,2,'a', 'b']
+    data3 = {'a':1, 2: 3}
+    # test without
+    assert Chepy(data1).without('ll').o == b'heo'
+    assert Chepy(data1).without('l', b'l').o == b'heo'
+    assert Chepy(data2).without(1, 'a').o == [2,'b']
+    assert Chepy(data3).without('a').o == {2: 3}
+    # test pick
+    assert Chepy(data1).pick('ll').o == b'll'
+    assert Chepy(data1).pick('l', b'l').o == b'll'
+    assert Chepy(data2).pick(1, 'a').o == [1,'a']
+    assert Chepy(data3).pick('a').o == {'a': 1}