diff --git a/text_wrapper.py b/text_wrapper.py new file mode 100644 index 0000000..4c5eeaf --- /dev/null +++ b/text_wrapper.py @@ -0,0 +1,127 @@ +"""Helper for inserting and extracting text from a string""" + +from typing import Optional + + +class TextWrapper: + """ + A class for inserting and extracting text from a string + """ + + def __init__(self, prefix: str, suffix: str) -> None: + """ + A class for inserting and extracting text from a string. + + The text is wrapped with a prefix and suffix, which are used to identify the + content across subsequent operations. + + Parameters: + - prefix: The prefix to wrap the content with + - suffix: The suffix to wrap the content with + """ + self._field_prefix = prefix + self._field_suffix = suffix + + def contains(self, text: str) -> bool: + """ + Check if the text contains the wrapped content + + Parameters: + - text: The text to check for the wrapped content + + Returns: + True if the text contains wrapped content, False otherwise + + >>> w = TextWrapper("=== START ===", "=== END ===") + >>> w.contains("This is some\\n=== START ===\\nwrapped\\n=== END ===\\ncontent") + True + >>> w.contains("This is some\\nunwrapped\\ncontent") + False + """ + return self._field_prefix in text and self._field_suffix in text + + def get(self, text: str) -> Optional[str]: + """ + Get the wrapped content from the provided text + + Parameters: + - text: The text to extract the wrapped content from + + Returns: + The wrapped content, or None if the text does not contain the + wrapped content + + >>> w = TextWrapper("=== START ===", "=== END ===") + >>> w.get("This is some\\n=== START ===\\nwrapped\\n=== END ===\\ncontent") + 'wrapped' + >>> w.get("This is some\\nunwrapped\\ncontent") # returns None + """ + start = text.find(self._field_prefix) + if start == -1: + return None + start += len(self._field_prefix) + end = text.find(self._field_suffix, start) + return text[start:end].strip() + + def remove(self, text: str) -> str: + """ + Remove the wrapped content from the text + + Parameters: + - text: The text to remove the wrapped content from + + Returns: + The text with the wrapped content removed + + >>> w = TextWrapper("=== START ===", "=== END ===") + >>> w.remove("This is some\\n=== START ===\\nwrapped\\n=== END ===\\ncontent") + 'This is some\\ncontent' + >>> w.remove("This is some\\nunwrapped\\ncontent") + 'This is some\\nunwrapped\\ncontent' + """ + start = text.find(self._field_prefix) + if start == -1: + return text + end = text.find(self._field_suffix, start) + front = text[:start] + back = text[end + len(self._field_suffix) :] + return front + back.lstrip() + + def upsert(self, text: str, new_content: str) -> str: + """ + Insert or update the wrapped content within the text + + Parameters: + - text: The text to insert or update the wrapped content in + - new_content: The new content to insert or update + + Returns: + The text with the wrapped content inserted or updated + + >>> w = TextWrapper("=== START ===", "=== END ===") + >>> w.upsert("This is some\\n=== START ===\\nwrapped\\n=== END ===\\ncontent", "new") + 'This is some\\n=== START ===\\nnew\\n=== END ===\\ncontent' + >>> w.upsert("This is some\\nunwrapped\\ncontent", "new") + 'This is some\\nunwrapped\\ncontent\\n=== START ===\\nnew\\n=== END ===' + """ + start = text.find(self._field_prefix) + if start == -1: + return ( + text + + "\n" + + self._field_prefix + + "\n" + + new_content + + "\n" + + self._field_suffix + ) + end = text.find(self._field_suffix, start) + return ( + text[:start] + + self._field_prefix + + "\n" + + new_content + + "\n" + + self._field_suffix + + text[end + len(self._field_suffix) :] + )