diff --git a/chepy/modules/utils.py b/chepy/modules/utils.py index 464e64e..ce3edd1 100644 --- a/chepy/modules/utils.py +++ b/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 diff --git a/chepy/modules/utils.pyi b/chepy/modules/utils.pyi index f80d32e..c00ec4c 100644 --- a/chepy/modules/utils.pyi +++ b/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: ... diff --git a/tests/test_utils.py b/tests/test_utils.py index 539a4af..b94b29b 100644 --- a/tests/test_utils.py +++ b/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} \ No newline at end of file