Make standard string methods return TextHandler again instead of str

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,