""" Clean a DataFrame column containing text data. """ import re import string from functools import partial, update_wrapper from typing import Any, Callable, Dict, List, Optional, Set, Union from unicodedata import normalize import dask.dataframe as dd import numpy as np import pandas as pd from ..assets.english_stopwords import english_stopwords from .utils import NULL_VALUES, to_dask REGEX_BRACKETS = { "angle": re.compile(r"(\<)[^<>]*(\>)"), "curly": re.compile(r"(\{)[^{}]*(\})"), "round": re.compile(r"(\()[^()]*(\))"), "square": re.compile(r"(\[)[^\[\]]*(\])"), } REGEX_DIGITS = re.compile(r"\d+") REGEX_DIGITS_BLOCK = re.compile(r"\b\d+\b") REGEX_HTML = re.compile(r"<[A-Za-z/][^>]*>|&(?:[a-z0-9]+|#[0-9]{1,6}|#x[0-9a-f]{1,6});") REGEX_PUNCTUATION = re.compile(fr"[{re.escape(string.punctuation)}]") REGEX_URL = re.compile(r"(?:https?://|www\.)\S+") REGEX_WHITESPACE = re.compile(r"[\n\t]|[ ]{2,}") [docs]def clean_text( df: Union[pd.DataFrame, dd.DataFrame], column: str, pipeline: Optional[List[Dict[str, Any]]] = None, stopwords: Optional[Set[str]] = None, ) -> pd.DataFrame: """ Clean text data in a DataFrame column. Read more in the :ref:`User Guide <clean_text_user_guide>`. Parameters ---------- df A pandas or Dask DataFrame containing the data to be cleaned. column The name of the column containing text data. pipeline A list of cleaning functions to be applied to the column. If None, use the default pipeline. See the :ref:`User Guide <clean_text_custom_pipeline>` for more information on customizing the pipeline. (default: None) stopwords A set of words to be removed from the column. If None, use NLTK's stopwords. (default: None) Examples -------- Clean a column of text data using the default pipeline. >>> df = pd.DataFrame({"text": ["This show was an amazing, fresh & innovative idea in the \ 70's when it first aired."]}) >>> clean_text(df, 'text') text 0 show amazing fresh innovative idea first aired """ df = to_dask(df) pipe = _get_default_pipeline(stopwords) if not pipeline else _get_custom_pipeline(pipeline) for func in pipe: df[column] = df[column].apply(func, meta=object) df = df.compute() return df [docs]def default_text_pipeline() -> List[Dict[str, Any]]: """ Return a list of dictionaries representing the functions in the default pipeline. Use as a template for creating a custom pipeline. Read more in the :ref:`User Guide <clean_text_user_guide>`. Examples -------- >>> default_text_pipeline() [{'operator': 'fillna'}, {'operator': 'lowercase'}, {'operator': 'remove_digits'}, {'operator': 'remove_html'}, {'operator': 'remove_urls'}, {'operator': 'remove_punctuation'}, {'operator': 'remove_accents'}, {'operator': 'remove_stopwords', 'parameters': {'stopwords': None}}, {'operator': 'remove_whitespace'}] """ return [ {"operator": "fillna"}, {"operator": "lowercase"}, {"operator": "remove_digits"}, {"operator": "remove_html"}, {"operator": "remove_urls"}, {"operator": "remove_punctuation"}, {"operator": "remove_accents"}, {"operator": "remove_stopwords", "parameters": {"stopwords": None}}, {"operator": "remove_whitespace"}, ] def _get_default_pipeline( stopwords: Optional[Set[str]] = None, ) -> List[Callable[..., Any]]: """ Return a list of functions defining the default pipeline. """ return [ _fillna, _lowercase, _remove_digits, _remove_html, _remove_urls, _remove_punctuation, _remove_accents, lambda x: _remove_stopwords(x, stopwords), _remove_whitespace, ] def _get_custom_pipeline(pipeline: List[Dict[str, Any]]) -> List[Callable[..., Any]]: """ Return a list of functions defining a custom pipeline. """ func_dict = _get_func_dict() custom_pipeline: List[Callable[..., Any]] = [] for component in pipeline: # Check whether function is built in or user defined operator = ( func_dict[component["operator"]] if isinstance(component["operator"], str) else component["operator"] ) # Append the function to the pipeline # If parameters are specified, create a partial function to lock in # the values and prevent them from being overwritten in subsequent loops if "parameters" in component: custom_pipeline.append(_wrapped_partial(operator, component["parameters"])) else: custom_pipeline.append(operator) return custom_pipeline def _get_func_dict() -> Dict[str, Callable[..., Any]]: """ Return a mapping of strings to function names. """ return { "fillna": _fillna, "lowercase": _lowercase, "sentence_case": _sentence_case, "title_case": _title_case, "uppercase": _uppercase, "remove_accents": _remove_accents, "remove_bracketed": _remove_bracketed, "remove_digits": _remove_digits, "remove_html": _remove_html, "remove_prefixed": _remove_prefixed, "remove_punctuation": _remove_punctuation, "remove_stopwords": _remove_stopwords, "remove_urls": _remove_urls, "remove_whitespace": _remove_whitespace, "replace_bracketed": _replace_bracketed, "replace_digits": _replace_digits, "replace_prefixed": _replace_prefixed, "replace_punctuation": _replace_punctuation, "replace_stopwords": _replace_stopwords, "replace_text": _replace_text, "replace_urls": _replace_urls, } def _fillna(text: Any, value: Any = np.nan) -> Any: """ Replace all null values with NaN (default) or the supplied value. """ return value if text in NULL_VALUES else str(text) def _lowercase(text: Any) -> Any: """ Convert all characters to lowercase. """ return str(text).lower() if pd.notna(text) else text def _sentence_case(text: Any) -> Any: """ Convert first character to uppercase and remaining to lowercase. """ return str(text).capitalize() if pd.notna(text) else text def _title_case(text: Any) -> Any: """ Convert first character of each word to uppercase and remaining to lowercase. """ return str(text).title() if pd.notna(text) else text def _uppercase(text: Any) -> Any: """ Convert all characters to uppercase. """ return str(text).upper() if pd.notna(text) else text def _remove_accents(text: Any) -> Any: """ Remove accents (diacritic marks). """ return ( normalize("NFD", str(text)).encode("ascii", "ignore").decode("ascii") if pd.notna(text) else text ) def _remove_bracketed(text: Any, brackets: Union[str, Set[str]], inclusive: bool = True) -> Any: """ Remove text between brackets. Parameters ---------- brackets The bracket style. - "angle": <> - "curly": {} - "round": () - "square": [] inclusive If True (default), remove the brackets along with the text in between. Otherwise, keep the brackets. """ if pd.isna(text): return text text = str(text) value = "" if inclusive else r"\g<1>\g<2>" if isinstance(brackets, set): for bracket in brackets: text = re.sub(REGEX_BRACKETS[bracket], value, text) else: text = re.sub(REGEX_BRACKETS[brackets], value, text) return text def _remove_digits(text: Any) -> Any: """ Remove all digits. """ return re.sub(REGEX_DIGITS, "", str(text)) if pd.notna(text) else text def _remove_html(text: Any) -> Any: """ Remove HTML tags. """ return re.sub(REGEX_HTML, "", str(text)) if pd.notna(text) else text def _remove_prefixed(text: Any, prefix: Union[str, Set[str]]) -> Any: """ Remove substrings that start with the prefix(es). """ if pd.isna(text): return text text = str(text) if isinstance(prefix, set): for pre in prefix: text = re.sub(fr"{pre}\S+", "", text) else: text = re.sub(fr"{prefix}\S+", "", text) return text def _remove_punctuation(text: Any) -> Any: """ Remove punctuation marks. """ return re.sub(REGEX_PUNCTUATION, " ", str(text)) if pd.notna(text) else text def _remove_stopwords(text: Any, stopwords: Optional[Set[str]] = None) -> Any: """ Remove a set of words from the text. If `stopwords` is None (default), use NLTK's stopwords. """ if pd.isna(text): return text stopwords = english_stopwords if not stopwords else stopwords return " ".join(word for word in str(text).split() if word.lower() not in stopwords) def _remove_urls(text: Any) -> Any: """ Remove URLS. """ return re.sub(REGEX_URL, "", str(text)) if pd.notna(text) else text def _remove_whitespace(text: Any) -> Any: """ Remove extra spaces along with tabs and newlines. """ return re.sub(REGEX_WHITESPACE, " ", str(text)).strip() if pd.notna(text) else text def _replace_bracketed( text: Any, brackets: Union[str, Set[str]], value: str, inclusive: bool = True ) -> Any: """ Replace text between brackets with the value. Parameters ---------- brackets The bracket style. - "angle": <> - "curly": {} - "round": () - "square": [] value The value to replace the text between the brackets. inclusive If True (default), replace the brackets with the new text as well. Otherwise, keep the brackets. """ if pd.isna(text): return text text = str(text) value = value if inclusive else fr"\g<1>{value}\g<2>" if isinstance(brackets, set): for bracket in brackets: text = re.sub(REGEX_BRACKETS[bracket], value, text) else: text = re.sub(REGEX_BRACKETS[brackets], value, text) return text def _replace_digits(text: Any, value: str, block: Optional[bool] = True) -> Any: """ Replace all digits with the value. If `block` is True (default), only replace blocks of digits. """ if pd.isna(text): return text return ( re.sub(REGEX_DIGITS_BLOCK, value, str(text)) if block else re.sub(REGEX_DIGITS, value, str(text)) ) def _replace_prefixed(text: Any, prefix: Union[str, Set[str]], value: str) -> Any: """ Replace all substrings starting with the prefix(es) with the value. """ if pd.isna(text): return text text = str(text) if isinstance(prefix, set): for pre in prefix: text = re.sub(fr"{pre}\S+", value, text) else: text = re.sub(fr"{prefix}\S+", value, text) return text def _replace_punctuation(text: Any, value: str) -> Any: """ Replace all punctuation marks with the value. """ return re.sub(REGEX_PUNCTUATION, value, str(text)) if pd.notna(text) else text def _replace_stopwords(text: Any, value: str, stopwords: Optional[Set[str]] = None) -> Any: """ Replace a set of words in the text with the value. If `stopwords` is None (default), use NLTK's stopwords. """ if pd.isna(text): return text stopwords = english_stopwords if not stopwords else stopwords return " ".join(word if word.lower() not in stopwords else value for word in str(text).split()) def _replace_text(text: Any, value: Dict[str, str], block: Optional[bool] = True) -> Any: """ Replace a sequence of characters with another according to the value mapping. If `block` is True (default), only replace standalone blocks of the sequence. """ if pd.isna(text): return text text = str(text) for old_value, new_value in value.items(): text = ( re.sub(fr"\b{old_value}\b", new_value, text, flags=re.IGNORECASE) if block else re.sub(fr"{old_value}", new_value, text, flags=re.IGNORECASE) ) return text def _replace_urls(text: Any, value: str) -> Any: """ Replace all URLs with the value. """ return re.sub(REGEX_URL, value, str(text)) if pd.notna(text) else text def _wrapped_partial( func: Callable[..., Callable[..., Any]], params: Dict[str, Any] ) -> Callable[..., Callable[..., Any]]: """ Return a partial function with a name and a doc attribute. """ partial_func = partial(func, **params) update_wrapper(partial_func, func) return partial_func