Concat removed prefix

explosion · Sep 26, 2023 · dff38bd · dff38bd
1 parent b2a974e
commit dff38bd
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/curated_transformers/util/string.py b/curated_transformers/util/string.py
@@ -135,7 +135,7 @@ def _revert(self, string: str) -> str:
             return string
 
 
-class StringRemovePrefix(StringSubRegEx):
+class StringRemovePrefix(StringTransform):
     """
     Strips a prefix from a given string.
     """
@@ -147,9 +147,14 @@ def __init__(self, prefix: str, *, reversible: bool = True):
         :param prefix:
             Prefix to be stripped.
         """
+
+        super().__init__(reversible)
+        self.prefix = prefix
+
+    def _apply(self, string: str) -> str:
         # TODO: Should be replaced with `removeprefix` once
         # Python 3.9 is the minimum requirement.
-        super().__init__(
-            forward=(f"^{re.escape(prefix)}", ""),
-            backward=(r"^(.)", f"{prefix}\\1") if reversible else None,
-        )
+        return re.sub(f"^{re.escape(self.prefix)}", "", string)
+
+    def _revert(self, string: str) -> str:
+        return f"{self.prefix}{string}"