Fix lstrip

Author: Timmy Welch

src/pyinfra/api/maskstring.py

@@ -198,6 +198,7 @@ def index(self, sub, start=0, end=sys.maxsize) -> int:
             return self.unmask().index(sub.unmask(), start, end)
         return super().index(sub, start, end)
 
+    # Not sure if this should return the unmasked result or the masked result
     # @override
     # def isalpha(self) -> bool:
     #     return self.unmask().isalpha()
@@ -265,13 +266,14 @@ def lower(self) -> "MaskString":
 
     @override
     def lstrip(self, chars=None) -> "MaskString":
-        s_chars = chars
+        chars_s = chars
         if isinstance(chars, MaskString):
-            s_chars = chars.unmask()
-            # explicitly set chars to the empty string as we want the resulting masked string to stay the same
-            chars = ""
-        return MaskString(self.unmask().lstrip(s_chars), masked_value=super().lstrip(chars))
+            # we want the resulting masked string to stay the same as otherwise it would most likely just strip all the characters
+            chars_s = chars.unmask()
+            return MaskString(self.unmask().lstrip(chars_s), masked_value=str(self))
+        return MaskString(self.unmask().lstrip(chars), masked_value=super().lstrip(chars))
 
+    # I don't know how maketrans works... default to using the masked value
     # maketrans = str.maketrans
 
     @override
@@ -337,11 +339,12 @@ def rpartition(self, sep) -> tuple["MaskString", "MaskString", "MaskString"]:
     def rstrip(self, chars=None) -> "MaskString":
         chars_s = chars
         if isinstance(chars, MaskString):
-            chars = chars.unmask()
-            # explicitly set chars to the empty string as we want the resulting masked string to stay the same
-            chars = ""
+            # we want the resulting masked string to stay the same as otherwise it would most likely just strip all the characters
+            chars_s = chars.unmask()
+            return MaskString(self.unmask().rstrip(chars_s), masked_value=str(self))
         return MaskString(self.unmask().rstrip(chars_s), masked_value=super().rstrip(chars))
 
+    # There is no way to keep the split functions "in-sync" with the masked string and the raw_value
     # @override
     # def split(self, sep=None, maxsplit=-1) -> list[MaskString]:
     #     if isinstance(sep, MaskString):
@@ -371,9 +374,9 @@ def startswith(self, prefix, start=0, end=sys.maxsize) -> bool:
     def strip(self, chars=None) -> "MaskString":
         chars_s = chars
         if isinstance(chars, MaskString):
+            # we want the resulting masked string to stay the same as otherwise it would most likely just strip all the characters
             chars_s = chars.unmask()
-            # explicitly set chars to the empty string as we want the resulting masked string to stay the same
-            chars = ""
+            return MaskString(self.unmask().strip(chars_s), masked_value=str(self))
         return MaskString(self.unmask().strip(chars_s), masked_value=super().strip(chars))
 
     @override
@@ -384,6 +387,7 @@ def swapcase(self) -> "MaskString":
     def title(self) -> "MaskString":
         return MaskString(self.unmask().title(), masked_value=super().title())
 
+    # I don't know how translate works... default to using the masked value
     # @override
     # def translate(self, *args) -> 'MaskString':
     #     return MaskString(self.unmask().translate(*args))

tests/test_api/test_api_maskstring.py

@@ -256,6 +256,11 @@ def test_lstrip(self):
         assert new_s.unmask() == "secret**"
         assert new_s == "*MASKED*"
 
+        s = MaskString("  secret  ")
+        new_s = s.lstrip(MaskString(" "))
+        assert new_s.unmask() == "secret  "
+        assert new_s == "*MASKED*"
+
     def test_partition(self):
         s = MaskString("secret:value")
         before, sep, after = s.partition(":")
@@ -448,6 +453,11 @@ def test_rstrip(self):
         assert new_s.unmask() == "**secret"
         assert new_s == "*MASKED*"
 
+        s = MaskString("  secret  ")
+        new_s = s.rstrip(MaskString(" "))
+        assert new_s.unmask() == "  secret"
+        assert new_s == "*MASKED*"
+
     def test_startswith(self):
         s = MaskString("secret value")
         assert s.startswith("*")
@@ -478,6 +488,11 @@ def test_strip(self):
         assert new_s.unmask() == "secret"
         assert new_s == "*MASKED*"
 
+        s = MaskString("  secret  ")
+        new_s = s.strip(MaskString(" "))
+        assert new_s.unmask() == "secret"
+        assert new_s == "*MASKED*"
+
     def test_swapcase(self):
         s = MaskString("Secret")
         new_s = s.swapcase()