diff --git a/scrapling/core/custom_types.py b/scrapling/core/custom_types.py index b8cb44f..f952da7 100644 --- a/scrapling/core/custom_types.py +++ b/scrapling/core/custom_types.py @@ -14,11 +14,70 @@ class TextHandler(str): __slots__ = () def __new__(cls, string): - # Because str is immutable and we can't override __init__ - if type(string) is str: + if isinstance(string, str): return super().__new__(cls, string) - else: - return super().__new__(cls, '') + return super().__new__(cls, '') + + # Make methods from original `str` class return `TextHandler` instead of returning `str` again + # Of course, this stupid workaround is only so we can keep the auto-completion working without issues in your IDE + # and I made sonnet write it for me :) + def strip(self, chars=None): + return TextHandler(super().strip(chars)) + + def lstrip(self, chars=None): + return TextHandler(super().lstrip(chars)) + + def rstrip(self, chars=None): + return TextHandler(super().rstrip(chars)) + + def capitalize(self): + return TextHandler(super().capitalize()) + + def casefold(self): + return TextHandler(super().casefold()) + + def center(self, width, fillchar=' '): + return TextHandler(super().center(width, fillchar)) + + def expandtabs(self, tabsize=8): + return TextHandler(super().expandtabs(tabsize)) + + def format(self, *args, **kwargs): + return TextHandler(super().format(*args, **kwargs)) + + def format_map(self, mapping): + return TextHandler(super().format_map(mapping)) + + def join(self, iterable): + return TextHandler(super().join(iterable)) + + def ljust(self, width, fillchar=' '): + return TextHandler(super().ljust(width, fillchar)) + + def rjust(self, width, fillchar=' '): + return TextHandler(super().rjust(width, fillchar)) + + def swapcase(self): + return TextHandler(super().swapcase()) + + def title(self): + return TextHandler(super().title()) + + def translate(self, table): + return TextHandler(super().translate(table)) + + def zfill(self, width): + return TextHandler(super().zfill(width)) + + def replace(self, old, new, count=-1): + return TextHandler(super().replace(old, new, count)) + + def upper(self): + return TextHandler(super().upper()) + + def lower(self): + return TextHandler(super().lower()) + ############## def sort(self, reverse: bool = False) -> str: """Return a sorted version of the string""" @@ -32,9 +91,9 @@ class TextHandler(str): def json(self) -> Dict: """Return json response if the response is jsonable otherwise throw error""" - # Using __str__ function as a workaround for orjson issue with subclasses of str + # Using str function as a workaround for orjson issue with subclasses of str # Check this out: https://github.com/ijl/orjson/issues/445 - return loads(self.__str__()) + return loads(str(self)) def re( self, regex: Union[str, Pattern[str]], replace_entities: bool = True, clean_match: bool = False,