diff --git a/docs/user_guide/text/TextFeatures.rst b/docs/user_guide/text/TextFeatures.rst index cf31dd0b9..1ad66bd4b 100644 --- a/docs/user_guide/text/TextFeatures.rst +++ b/docs/user_guide/text/TextFeatures.rst @@ -38,27 +38,30 @@ Text features :class:`TextFeatures()` can extract the following features from a text piece: -- **char_count**: Number of characters in the text +- **char_count**: Number of characters, excluding whitespace - **word_count**: Number of words (whitespace-separated tokens) - **sentence_count**: Number of sentences (based on .!? punctuation) -- **avg_word_length**: Average length of words +- **avg_word_length**: Average number of characters per word - **digit_count**: Number of digit characters -- **letter_count**: Number of alphabetic characters (a-z, A-Z) -- **uppercase_count**: Number of uppercase letters -- **lowercase_count**: Number of lowercase letters -- **special_char_count**: Number of special characters (non-alphanumeric) +- **letter_count**: Number of letters a-z and A-Z +- **uppercase_count**: Number of uppercase letters A-Z +- **lowercase_count**: Number of lowercase letters a-z +- **special_char_count**: Number of characters that are not a-z, A-Z, 0-9 or whitespace - **whitespace_count**: Number of whitespace characters - **whitespace_ratio**: Ratio of whitespace to total characters -- **digit_ratio**: Ratio of digits to total characters -- **uppercase_ratio**: Ratio of uppercase to total characters +- **digit_ratio**: Ratio of digits to non-whitespace characters +- **uppercase_ratio**: Ratio of uppercase letters to non-whitespace characters - **has_digits**: Binary indicator if text contains digits -- **has_uppercase**: Binary indicator if text contains uppercase +- **has_uppercase**: Binary indicator if text contains uppercase letters A-Z - **is_empty**: Binary indicator if text is empty -- **starts_with_uppercase**: Binary indicator if text starts with uppercase +- **starts_with_uppercase**: Binary indicator if text starts with A-Z - **ends_with_punctuation**: Binary indicator if text ends with .!? - **unique_word_count**: Number of unique words (case-insensitive) - **lexical_diversity**: Ratio of unique words to total words +Letters with accents or from other alphabets, like é or ß, are not counted as letters +or uppercase letters; they are counted as special characters. + The **number of sentences** is inferred by :class:`TextFeatures()` by counting blocks of sentence-ending punctuation (., !, ?) as a proxy for sentence boundaries. This means that multiple consecutive punctuation marks (e.g., "!!!" or "??") are counted as a single @@ -160,7 +163,7 @@ The input dataframe looks like this: Now let's extract 5 specific text features: the number of words, the number of characters, the number of sentences, whether the text has digits, and the ratio of -upper- to lowercase: +uppercase letters to non-whitespace characters: .. code:: python @@ -221,10 +224,10 @@ The output dataframe contains all 20 text features extracted from the `review` c 3 TERRIBLE!!! DO NOT BUY! Awful 20 4 review_sentence_count review_avg_word_length review_digit_count review_letter_count - 0 2 6.285714 0 36 - 1 2 6.200000 0 25 - 2 2 3.888889 2 23 - 3 2 5.750000 0 16 + 0 2 5.428571 0 36 + 1 2 5.400000 0 25 + 2 2 3.000000 2 23 + 3 2 5.000000 0 16 review_uppercase_count review_lowercase_count review_special_char_count review_whitespace_count 0 9 27 2 6 @@ -279,6 +282,57 @@ extracted features remain: 2 Average 9 27 3 Awful 4 20 +With polars +~~~~~~~~~~~ + +:class:`TextFeatures()` works the same way with a polars dataframe, and returns a polars +dataframe. Let's create a toy dataset with a missing value: + +.. code:: python + + import polars as pl + from feature_engine.text import TextFeatures + + X = pl.DataFrame({ + 'review': [ + 'This product is AMAZING! Best purchase ever.', + 'Not great. Would not recommend.', + 'OK for the price. 3 out of 5 stars.', + None, + ], + }) + +Let's extract the number of words, whether the text has digits, and the ratio of +uppercase letters: + +.. code:: python + + tf = TextFeatures( + variables=['review'], + features=['word_count', 'has_digits', 'uppercase_ratio'], + ) + + X_transformed = tf.fit_transform(X) + + print(X_transformed) + +We obtain a polars dataframe with the new features. The missing value was replaced by +an empty string, which has 0 words: + +.. code-block:: none + + shape: (4, 4) + ┌─────────────────────────────────┬───────────────────┬───────────────────┬────────────────────────┐ + │ review ┆ review_word_count ┆ review_has_digits ┆ review_uppercase_ratio │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ str ┆ i64 ┆ i64 ┆ f64 │ + ╞═════════════════════════════════╪═══════════════════╪═══════════════════╪════════════════════════╡ + │ This product is AMAZING! Best … ┆ 7 ┆ 0 ┆ 0.236842 │ + │ Not great. Would not recommend… ┆ 5 ┆ 0 ┆ 0.074074 │ + │ OK for the price. 3 out of 5 s… ┆ 9 ┆ 1 ┆ 0.074074 │ + │ ┆ 0 ┆ 0 ┆ 0.0 │ + └─────────────────────────────────┴───────────────────┴───────────────────┴────────────────────────┘ + Combining with sklearn's bag-of-words ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/feature_engine/text/text_features.py b/feature_engine/text/text_features.py index 5198639d3..3d492edf6 100644 --- a/feature_engine/text/text_features.py +++ b/feature_engine/text/text_features.py @@ -1,8 +1,12 @@ # Authors: Ankit Hemant Lade (contributor) # License: BSD 3 clause +import string +from functools import cached_property from typing import List, Optional, Union, cast -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -17,35 +21,170 @@ check_X, ) -# Available text features and their computation functions +# The characters Python treats as whitespace. Listed explicitly because the \s of +# polars' regex engine misses \x1c-\x1f, and we want the same counts everywhere. +_WHITESPACE = ( + "\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\xa0\u1680\u2000\u2001\u2002\u2003\u2004" + "\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000" +) +_WORD = f"[^{_WHITESPACE}]+" + +# Each feature is computed from the text statistics of one of the classes below, +# so all backends share the same definitions. TEXT_FEATURES = { - "char_count": lambda x: x.str.replace(r"\s+", "", regex=True).str.len(), - "word_count": lambda x: x.str.strip().str.split().str.len(), - "sentence_count": lambda x: x.str.count(r"[.!?]+"), - "avg_word_length": lambda x: x.str.strip().str.len() - / x.str.strip().str.split().str.len(), - "digit_count": lambda x: x.str.count(r"\d"), - "letter_count": lambda x: x.str.count(r"[a-zA-Z]"), - "uppercase_count": lambda x: x.str.count(r"[A-Z]"), - "lowercase_count": lambda x: x.str.count(r"[a-z]"), - "special_char_count": lambda x: x.str.count(r"[^a-zA-Z0-9\s]"), - "whitespace_count": lambda x: x.str.count(r"\s"), - "whitespace_ratio": lambda x: x.str.count(r"\s") / x.str.len().replace(0, 1), - "digit_ratio": lambda x: x.str.count(r"\d") - / x.str.replace(r"\s+", "", regex=True).str.len().replace(0, 1), - "uppercase_ratio": lambda x: x.str.count(r"[A-Z]") - / x.str.replace(r"\s+", "", regex=True).str.len().replace(0, 1), - "has_digits": lambda x: x.str.contains(r"\d", regex=True).astype(int), - "has_uppercase": lambda x: x.str.contains(r"[A-Z]", regex=True).astype(int), - "is_empty": lambda x: (x.str.len() == 0).astype(int), - "starts_with_uppercase": lambda x: x.str.match(r"^[A-Z]").astype(int), - "ends_with_punctuation": lambda x: x.str.match(r".*[.!?]$").astype(int), - "unique_word_count": lambda x: (x.str.lower().str.split().apply(set).str.len()), - "lexical_diversity": lambda x: x.str.lower().str.split().apply(set).str.len() - / x.str.strip().str.split().str.len(), + "char_count": lambda t: t.length - t.count_in(_WHITESPACE), + "word_count": lambda t: t.word_count, + "sentence_count": lambda t: t.count(r"[.!?]+"), + "avg_word_length": lambda t: (t.length - t.count_in(_WHITESPACE)) + / t.word_count.clip(1), + "digit_count": lambda t: t.count(r"\d"), + "letter_count": lambda t: t.count_in(string.ascii_letters), + "uppercase_count": lambda t: t.count(r"[A-Z]"), + "lowercase_count": lambda t: t.count_in(string.ascii_lowercase), + "special_char_count": lambda t: t.count_not_in( + string.ascii_letters + string.digits + _WHITESPACE + ), + "whitespace_count": lambda t: t.count_in(_WHITESPACE), + "whitespace_ratio": lambda t: t.count_in(_WHITESPACE) / t.length.clip(1), + "digit_ratio": lambda t: t.count(r"\d") + / (t.length - t.count_in(_WHITESPACE)).clip(1), + "uppercase_ratio": lambda t: t.count(r"[A-Z]") + / (t.length - t.count_in(_WHITESPACE)).clip(1), + "has_digits": lambda t: t.contains(r"\d"), + "has_uppercase": lambda t: t.contains(r"[A-Z]"), + "is_empty": lambda t: t.is_empty(), + "starts_with_uppercase": lambda t: t.contains(r"^[A-Z]"), + "ends_with_punctuation": lambda t: t.ends_with_punctuation(), + "unique_word_count": lambda t: t.unique_word_count, + "lexical_diversity": lambda t: t.unique_word_count / t.word_count.clip(1), } +class _PandasText: + """Text statistics of a pandas Series of strings.""" + + def __init__(self, text, native_namespace): + self.text = text + self._pd = native_namespace + # several features share the same counts, and pandas computes them eagerly + self._counts: dict = {} + + @cached_property + def length(self): + return self.text.str.len() + + @cached_property + def word_count(self): + # a Python loop is 2x faster than pandas' str.split().str.len() + words = [len(s.split()) for s in self.text.tolist()] + return self._pd.Series(words, index=self.text.index) + + @cached_property + def unique_word_count(self): + words = [len(set(s.lower().split())) for s in self.text.tolist()] + return self._pd.Series(words, index=self.text.index) + + def count(self, pattern): + if ("count", pattern) not in self._counts: + self._counts[("count", pattern)] = self.text.str.count(pattern) + return self._counts[("count", pattern)] + + def count_in(self, characters): + # deleting the characters with translate is faster than a regex count + if ("count_in", characters) not in self._counts: + self._counts[("count_in", characters)] = ( + self.length - self.count_not_in(characters) + ) + return self._counts[("count_in", characters)] + + def count_not_in(self, characters): + table = str.maketrans("", "", characters) + return self.text.str.translate(table).str.len() + + def contains(self, pattern): + return self.text.str.contains(pattern, regex=True).astype(int) + + def is_empty(self): + return self.text.eq("").astype(int) + + def ends_with_punctuation(self): + return self.text.str.match(r".*[.!?]$").astype(int) + + +class _NarwhalsText: + """Text statistics of a string column, as narwhals expressions.""" + + def __init__(self, text, namespace): + self.text = text + self._ns = namespace + + @property + def length(self): + return self.text.str.len_chars().cast(self._ns.Int64) + + @property + def word_count(self): + return self.count(_WORD) + + @property + def unique_word_count(self): + words = ( + self.text.str.to_lowercase() + .str.replace_all(f"[{_WHITESPACE}]+", " ") + .str.strip_chars(" ") + .str.split(" ") + ) + # splitting a text without words returns one empty word + return ( + self._ns.when(self.word_count == 0) + .then(0) + .otherwise(words.list.unique().list.len()) + .cast(self._ns.Int64) + ) + + def count(self, pattern): + # narwhals can't count matches, but replacing each match with 2 + # characters instead of 1 makes the text 1 character longer per match + return ( + self.text.str.replace_all(pattern, "ab").str.len_chars() + - self.text.str.replace_all(pattern, "a").str.len_chars() + ).cast(self._ns.Int64) + + def count_in(self, characters): + return self.length - self.count_not_in(characters) + + def count_not_in(self, characters): + kept = self.text.str.replace_all(f"[{characters}]", "") + return kept.str.len_chars().cast(self._ns.Int64) + + def contains(self, pattern): + return self.text.str.contains(pattern).cast(self._ns.Int64) + + def is_empty(self): + return (self.text.str.len_chars() == 0).cast(self._ns.Int64) + + def ends_with_punctuation(self): + # Python's regex $ also matches before a final \n, which would let + # "x.\n\n" match on backends that use it + ends = self.text.str.contains(r"^[^\n]*[.!?]\n?$") + return (ends & ~self.text.str.ends_with("\n\n")).cast(self._ns.Int64) + + +class _PolarsText(_NarwhalsText): + """Text statistics of a string column, as polars expressions.""" + + @property + def unique_word_count(self): + words = self.text.str.to_lowercase().str.extract_all(_WORD) + return words.list.n_unique().cast(self._ns.Int64) + + def count(self, pattern): + return self.text.str.count_matches(pattern).cast(self._ns.Int64) + + def count_not_in(self, characters): + return self.count(f"[^{characters}]") + + class TextFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): """ TextFeatures() extracts numerical features from text/string variables. This @@ -64,23 +203,25 @@ class TextFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): features: list, default=None List of text features to extract. Available features are: - - 'char_count': Number of characters in the text + - 'char_count': Number of characters, excluding whitespace - 'word_count': Number of words (whitespace-separated tokens) - 'sentence_count': Number of sentences (based on .!? punctuation) - - 'avg_word_length': Average length of words + - 'avg_word_length': Average number of characters per word - 'digit_count': Number of digit characters - - 'letter_count': Number of alphabetic characters (a-z, A-Z) - - 'uppercase_count': Number of uppercase letters - - 'lowercase_count': Number of lowercase letters - - 'special_char_count': Number of special characters (non-alphanumeric) + - 'letter_count': Number of letters a-z and A-Z + - 'uppercase_count': Number of uppercase letters A-Z + - 'lowercase_count': Number of lowercase letters a-z + - 'special_char_count': Number of characters that are not a-z, A-Z, 0-9 + or whitespace - 'whitespace_count': Number of whitespace characters - 'whitespace_ratio': Ratio of whitespace to total characters - - 'digit_ratio': Ratio of digits to total characters - - 'uppercase_ratio': Ratio of uppercase to total characters + - 'digit_ratio': Ratio of digits to non-whitespace characters + - 'uppercase_ratio': Ratio of uppercase letters to non-whitespace + characters - 'has_digits': Binary indicator if text contains digits - - 'has_uppercase': Binary indicator if text contains uppercase + - 'has_uppercase': Binary indicator if text contains uppercase letters A-Z - 'is_empty': Binary indicator if text is empty - - 'starts_with_uppercase': Binary indicator if text starts with uppercase + - 'starts_with_uppercase': Binary indicator if text starts with A-Z - 'ends_with_punctuation': Binary indicator if text ends with .!? - 'unique_word_count': Number of unique words (case-insensitive) - 'lexical_diversity': Ratio of unique words to total words @@ -88,9 +229,9 @@ class TextFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): If None, extracts all available features. missing_values: string, default='ignore' - If 'ignore', NaNs will be filled with an empty string before feature - extraction. If 'raise', the transformer will raise an error if missing data - is found. + If 'ignore', missing values will be filled with an empty string before + feature extraction. If 'raise', the transformer will raise an error if + missing data is found. drop_original: bool, default=False Whether to drop the original text columns after transformation. @@ -141,8 +282,6 @@ class TextFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): ... features=['char_count', 'word_count', 'has_digits'] ... ) >>> tf.fit(X) - TextFeatures(features=['char_count', 'word_count', 'has_digits'], - variables=['text']) >>> X = tf.transform(X) >>> pd.options.display.max_columns = 10 >>> print(X) @@ -160,7 +299,6 @@ def __init__( drop_original: bool = False, ) -> None: - # Validate variables if isinstance(variables, str): variables = [variables] if not isinstance(variables, list) or not all( @@ -168,24 +306,17 @@ def __init__( ): raise ValueError( "variables must be a string or a list of strings. " - f"Got {type(variables).__name__} instead." + f"Got {variables} instead." ) - # Validate features - if features is not None: - if not isinstance(features, list) or not all( - isinstance(f, str) for f in features - ): - raise ValueError( - "features must be None or a list of strings. " - f"Got {type(features).__name__} instead." - ) - invalid_features = set(features) - set(TEXT_FEATURES.keys()) - if invalid_features: - raise ValueError( - f"Invalid features: {invalid_features}. " - f"Available features are: {list(TEXT_FEATURES.keys())}" - ) + if features is not None and ( + not isinstance(features, list) + or not all(isinstance(f, str) and f in TEXT_FEATURES for f in features) + ): + raise ValueError( + "features must be None or a list with any of " + f"{list(TEXT_FEATURES.keys())}. Got {features} instead." + ) _check_param_drop_original(drop_original) _check_param_missing_values(missing_values) @@ -195,38 +326,38 @@ def __init__( self.missing_values = missing_values self.drop_original = drop_original - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y=None): """ This transformer does not learn any parameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just the variables to transform. - y: pandas Series, or np.array. Defaults to None. + y: Series, or np.array. Defaults to None. The target. It is not needed in this transformer. You can pass y or None. """ + nw_X = check_X(X) - # check input dataframe - X = check_X(X) - - # Validate user-specified variables exist - missing = set(self.variables) - set(X.columns) - if missing: + missing = set(self.variables) - set(nw_X.columns) + if len(missing) > 0: raise ValueError(f"Variables {missing} are not present in the dataframe.") - # Validate that the variables are object or string - non_text = [ - col - for col in self.variables - if not ( - pd.api.types.is_string_dtype(X[col]) - or pd.api.types.is_object_dtype(X[col]) - ) - ] - if non_text: + non_text = [] + for var in self.variables: + dtype = nw_X.get_column(var).dtype + # pandas categories can be numbers, polars categories are always strings + if isinstance(dtype, nw.Categorical) and nwd.is_pandas_dataframe(X) is True: + is_text = X[var].cat.categories.inferred_type == "string" + else: + is_text = isinstance( + dtype, (nw.String, nw.Object, nw.Categorical, nw.Enum) + ) + if is_text is False: + non_text.append(var) + if len(non_text) > 0: raise ValueError( f"Variables {non_text} are not object or string. " "Please provide text variables only." @@ -234,103 +365,99 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): self.variables_ = self.variables - # check if dataset contains na if self.missing_values == "raise": _check_contains_na( X, cast(list[Union[str, int]], self.variables_), error_msg="optional" ) - # Set features to extract if self.features is None: self.features_ = list(TEXT_FEATURES.keys()) else: self.features_ = self.features - # save input features - self.feature_names_in_ = X.columns.tolist() - - # save train set shape - self.n_features_in_ = X.shape[1] + self.feature_names_in_ = nw_X.columns + self.n_features_in_ = nw_X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Extract text features and add them to the dataframe. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas dataframe + X_new: dataframe The dataframe with the original columns plus the new text features. """ - - # Check method fit has been called check_is_fitted(self) + nw_X = check_X(X) + _check_X_matches_training_df(nw_X, self.n_features_in_) - # check that input is a dataframe - X = check_X(X) - - # Check if input data contains same number of columns as dataframe used to fit. - _check_X_matches_training_df(X, self.n_features_in_) - - # check if dataset contains na if self.missing_values == "raise": _check_contains_na( X, cast(list[Union[str, int]], self.variables_), error_msg="optional" ) - else: - X[self.variables_] = X[self.variables_].fillna("") - # reorder variables to match train set - X = X[self.feature_names_in_] - - # Extract features for each text variable - for var in self.variables_: - for feature_name in self.features_: - new_col_name = f"{var}_{feature_name}" - feature_func = TEXT_FEATURES[feature_name] - X[new_col_name] = feature_func(X[var]) - - # Fill any NaN values resulting from computation with 0 - X[new_col_name] = X[new_col_name].fillna(0) + # pandas is faster than narwhals. + if nwd.is_pandas_dataframe(X) is True: + X_new = self._transform_pandas(X, nw.get_native_namespace(nw_X)) + elif nwd.is_polars_dataframe(X) is True: + # polars counts regex matches natively, narwhals needs two replacements. + X_new = self._transform_expressions( + X, nw.get_native_namespace(nw_X), _PolarsText + ) + else: + X_new = self._transform_expressions(nw_X, nw, _NarwhalsText).to_native() - if self.drop_original: - X = X.drop(columns=self.variables_) + return X_new - return X + def _transform_pandas(self, X, native_namespace): + X_new = X[self.feature_names_in_] + if self.missing_values == "ignore": + X_new = X_new.fillna({var: "" for var in self.variables_}) - def get_feature_names_out(self, input_features=None) -> List[str]: - """ - Get output feature names for transformation. + new_features = [] + for var in self.variables_: + statistics = _PandasText(X_new[var], native_namespace) + new_features += [ + TEXT_FEATURES[feature](statistics).rename(f"{var}_{feature}") + for feature in self.features_ + ] - Parameters - ---------- - input_features : array-like of str or None, default=None - Input features. If ``None``, uses ``feature_names_in_``. + X_new = native_namespace.concat([X_new, *new_features], axis=1) + if self.drop_original is True: + X_new = X_new.drop(columns=self.variables_) - Returns - ------- - feature_names_out : list of str - Output feature names. - """ - check_is_fitted(self) + return X_new - # Start with original features - if self.drop_original: - feature_names = [ - f for f in self.feature_names_in_ if f not in self.variables_ + def _transform_expressions(self, X, namespace, text_class): + # polars and narwhals expressions share the API used here + filled_text, new_features = [], [] + for var in self.variables_: + if self.missing_values == "ignore": + filled_text.append(namespace.col(var).fill_null("")) + text = namespace.col(var).cast(namespace.String).fill_null("") + statistics = text_class(text, namespace) + new_features += [ + TEXT_FEATURES[feature](statistics).alias(f"{var}_{feature}") + for feature in self.features_ ] - else: - feature_names = list(self.feature_names_in_) - # Add new text feature names - for var in self.variables_: - for feature_name in self.features_: - feature_names.append(f"{var}_{feature_name}") + X_new = X.select(self.feature_names_in_).with_columns( + *filled_text, *new_features + ) + if self.drop_original is True: + X_new = X_new.drop(self.variables_) + + return X_new - return feature_names + def _get_new_features_name(self) -> List[str]: + """Return the names of the created features.""" + return [ + f"{var}_{feature}" for var in self.variables_ for feature in self.features_ + ] diff --git a/tests/test_text/test_text_features.py b/tests/test_text/test_text_features.py index 827ff6391..26568748f 100644 --- a/tests/test_text/test_text_features.py +++ b/tests/test_text/test_text_features.py @@ -1,798 +1,574 @@ +import re +from datetime import datetime +from types import SimpleNamespace + +import narwhals as nw +import numpy as np import pandas as pd +import polars as pl import pytest +from sklearn.exceptions import NotFittedError -from feature_engine.text import TextFeatures +from feature_engine.text import TextFeatures, text_features from feature_engine.text.text_features import TEXT_FEATURES - -# ============================================================================== -# INIT TESTS -# ============================================================================== +from tests.backend_helpers import frame_to_dict, make_series + +TEXT = [ + "Hello World!", + "HELLO", + "12345", + "e.g. i.e.", + " ", + " trailing ", + "abc...", + "", + None, + "A? B! C.", + "HeLLo", + "Hi! @#", + "A1b2 C3d4!@#$", + "???", + "i.e., this is wrong", + "Is 1 > 2? No, 100%!", + "Hello. World", + "Hello. World.", + "Hello... World!?!", + "This is a proper sentence containing " + "supercalifragilisticexpialidocious and exceptionally long words.", +] + +# non-ASCII letters and digits, non-breaking space (\xa0), file separator (\x1c), +# new lines around the final punctuation and a Greek final sigma +TEXT_EDGE_CASES = [ + "", + None, + " ", + "Hello World!", + "HELLO", + "\N{LATIN CAPITAL LETTER E WITH ACUTE}COLE " + "na\N{LATIN SMALL LETTER I WITH DIAERESIS}ve 123", + "\N{ARABIC-INDIC DIGIT THREE} digits", + "a\xa0b\x1cc", + "x.\n", + "x.\n\n", + "a\nb.", + "Dog dog DOG", + "\N{GREEK CAPITAL LETTER OMICRON}\N{GREEK CAPITAL LETTER DELTA}" + "\N{GREEK CAPITAL LETTER OMICRON}\N{GREEK CAPITAL LETTER SIGMA} " + "\N{GREEK SMALL LETTER OMICRON}\N{GREEK SMALL LETTER DELTA}" + "\N{GREEK SMALL LETTER OMICRON}\N{GREEK SMALL LETTER FINAL SIGMA}", + "Is 1 > 2? No, 100%!", +] + +EXPECTED = { + "char_count": [11, 5, 5, 8, 0, 8, 6, 0, 0, 6, 5, 5, 12, 3, 16, 14, 11, 12, 16, 91], + "word_count": [2, 1, 1, 2, 0, 1, 1, 0, 0, 3, 1, 2, 2, 1, 4, 6, 2, 2, 2, 11], + "sentence_count": [1, 0, 0, 4, 0, 0, 1, 0, 0, 3, 0, 1, 1, 1, 2, 2, 1, 2, 2, 1], + "avg_word_length": [ + 11 / 2, 5, 5, 4, 0, 8, 6, 0, 0, 2, 5, 5 / 2, 6, 3, 4, 14 / 6, 11 / 2, 6, 8, + 91 / 11, + ], + "digit_count": [0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 5, 0, 0, 0, 0], + "letter_count": [10, 5, 0, 4, 0, 8, 3, 0, 0, 3, 5, 2, 4, 0, 13, 4, 10, 10, 10, 90], + "uppercase_count": [2, 5, 0, 0, 0, 0, 0, 0, 0, 3, 3, 1, 2, 0, 0, 2, 2, 2, 2, 1], + "lowercase_count": [8, 0, 0, 4, 0, 8, 3, 0, 0, 0, 2, 1, 2, 0, 13, 2, 8, 8, 8, 89], + "special_char_count": [1, 0, 0, 4, 0, 0, 3, 0, 0, 3, 0, 3, 4, 3, 3, 5, 1, 2, 6, 1], + "whitespace_count": [1, 0, 0, 1, 3, 2, 0, 0, 0, 2, 0, 1, 1, 0, 3, 5, 1, 1, 1, 10], + "whitespace_ratio": [ + 1 / 12, 0, 0, 1 / 9, 1, 2 / 10, 0, 0, 0, 2 / 8, 0, 1 / 6, 1 / 13, 0, 3 / 19, + 5 / 19, 1 / 12, 1 / 13, 1 / 17, 10 / 101, + ], + "digit_ratio": [ + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4 / 12, 0, 0, 5 / 14, 0, 0, 0, 0, + ], + "uppercase_ratio": [ + 2 / 11, 1, 0, 0, 0, 0, 0, 0, 0, 3 / 6, 3 / 5, 1 / 5, 2 / 12, 0, 0, 2 / 14, + 2 / 11, 2 / 12, 2 / 16, 1 / 91, + ], + "has_digits": [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0], + "has_uppercase": [1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1], + "is_empty": [0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "starts_with_uppercase": [ + 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, + ], + "ends_with_punctuation": [ + 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, + ], + "unique_word_count": [2, 1, 1, 2, 0, 1, 1, 0, 0, 3, 1, 2, 2, 1, 4, 6, 2, 2, 2, 11], + "lexical_diversity": [1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], +} + +EXPECTED_EDGE_CASES = { + "char_count": [0, 0, 0, 11, 5, 13, 7, 3, 2, 2, 3, 9, 8, 14], + "word_count": [0, 0, 0, 2, 1, 3, 2, 3, 1, 1, 2, 3, 2, 6], + "sentence_count": [0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 2], + "avg_word_length": [ + 0, 0, 0, 11 / 2, 5, 13 / 3, 7 / 2, 1, 2, 2, 3 / 2, 3, 4, 14 / 6, + ], + "digit_count": [0, 0, 0, 0, 0, 3, 1, 0, 0, 0, 0, 0, 0, 5], + "letter_count": [0, 0, 0, 10, 5, 8, 6, 3, 1, 1, 2, 9, 0, 4], + "uppercase_count": [0, 0, 0, 2, 5, 4, 0, 0, 0, 0, 0, 4, 0, 2], + "lowercase_count": [0, 0, 0, 8, 0, 4, 6, 3, 1, 1, 2, 5, 0, 2], + "special_char_count": [0, 0, 0, 1, 0, 2, 1, 0, 1, 1, 1, 0, 8, 5], + "whitespace_count": [0, 0, 3, 1, 0, 2, 1, 2, 1, 2, 1, 2, 1, 5], + "whitespace_ratio": [ + 0, 0, 1, 1 / 12, 0, 2 / 15, 1 / 8, 2 / 5, 1 / 3, 2 / 4, 1 / 4, 2 / 11, 1 / 9, + 5 / 19, + ], + "digit_ratio": [0, 0, 0, 0, 0, 3 / 13, 1 / 7, 0, 0, 0, 0, 0, 0, 5 / 14], + "uppercase_ratio": [0, 0, 0, 2 / 11, 1, 4 / 13, 0, 0, 0, 0, 0, 4 / 9, 0, 2 / 14], + "has_digits": [0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1], + "has_uppercase": [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1], + "is_empty": [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "starts_with_uppercase": [0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1], + "ends_with_punctuation": [0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1], + "unique_word_count": [0, 0, 0, 2, 1, 3, 2, 3, 1, 1, 2, 1, 1, 6], + "lexical_diversity": [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1 / 3, 1 / 2, 1], +} + + +# init parameters +@pytest.mark.parametrize( + "variables", [123, True, None, [1, 2], ["text", 123], ("text",), {"text": 1}] +) +def test_error_if_variables_not_string_or_list_of_strings(variables): + msg = f"variables must be a string or a list of strings. Got {variables} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + TextFeatures(variables=variables) @pytest.mark.parametrize( - "invalid_variables", + "features", [ + "char_count", 123, True, + ("char_count",), + {"char_count": 1}, [1, 2], - ["text", 123], - {"text": 1}, + ["char_count", True], + ["invalid_feature"], + ["char_count", "invalid_feature"], ], ) -def test_invalid_variables_raises_error(invalid_variables): - with pytest.raises(ValueError, match="variables must be a string or a list of"): - TextFeatures(variables=invalid_variables) +def test_error_if_features_not_permitted(features): + msg = ( + f"features must be None or a list with any of {list(TEXT_FEATURES.keys())}. " + f"Got {features} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + TextFeatures(variables=["text"], features=features) + + +@pytest.mark.parametrize("missing_values", ["empanada", True, 1, None, ["raise"]]) +def test_error_if_missing_values_not_permitted(missing_values): + msg = ( + "missing_values takes only values 'raise' or 'ignore'. " + f"Got {missing_values} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + TextFeatures(variables=["text"], missing_values=missing_values) + + +@pytest.mark.parametrize("drop_original", ["True", 1, None, [True]]) +def test_error_if_drop_original_not_bool(drop_original): + msg = ( + "drop_original takes only boolean values True and False. " + f"Got {drop_original} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + TextFeatures(variables=["text"], drop_original=drop_original) @pytest.mark.parametrize( - "invalid_features, err_msg", + "features, missing_values, drop_original", [ - ("some_string", "features must be"), - ([1, 2], "features must be"), - (123, "features must be"), - (True, "features must be"), - (["some_string", True], "features must be"), - ({"some_string": 1}, "features must be"), - (["invalid_feature"], "Invalid features"), - (["char_count", "invalid_feature"], "Invalid features"), + (None, "ignore", False), + (["char_count"], "raise", True), + (["word_count", "lexical_diversity"], "ignore", True), ], ) -def test_invalid_features_raises_error(invalid_features, err_msg): - with pytest.raises(ValueError, match=err_msg): - TextFeatures(variables=["text"], features=invalid_features) - - -# ============================================================================== -# FIT TESTS -# ============================================================================== +def test_init_param_assignment(features, missing_values, drop_original): + transformer = TextFeatures( + variables=["text"], + features=features, + missing_values=missing_values, + drop_original=drop_original, + ) + assert transformer.features == features + assert transformer.missing_values == missing_values + assert transformer.drop_original is drop_original +# fit and transform @pytest.mark.parametrize( - "variables, features", + "variables, features, variables_, features_", [ - ("text", None), - (["string"], ["char_count"]), - (["text", "string"], ["sentence_count", "avg_word_length"]), + ("text", None, ["text"], list(TEXT_FEATURES.keys())), + (["string"], ["char_count"], ["string"], ["char_count"]), + (["text", "string"], ["word_count"], ["text", "string"], ["word_count"]), ], ) -def test_fit_stores_attributes(variables, features): - X = pd.DataFrame({"text": ["Hello"], "string": ["Bye"]}) - transformer = TextFeatures(variables=variables, features=features) - transformer.fit(X) - - assert ( - transformer.variables_ == variables - if isinstance(variables, list) - else transformer.variables_ == [variables] - ) - assert ( - transformer.features_ == list(TEXT_FEATURES.keys()) - if features is None - else transformer.features_ == features - ) - assert transformer.feature_names_in_ == ["text", "string"] - assert transformer.n_features_in_ == 2 +def test_fit_attributes(make_df, variables, features, variables_, features_): + X = make_df({"text": ["Hello"], "string": ["Bye"], "number": [1]}) + transformer = TextFeatures(variables=variables, features=features).fit(X) + + assert transformer.variables_ == variables_ + assert transformer.features_ == features_ + assert transformer.feature_names_in_ == ["text", "string", "number"] + assert transformer.n_features_in_ == 3 + + +@pytest.mark.parametrize("target", ["series", "list", "array"]) +def test_fit_ignores_the_target(make_df, target): + X = make_df({"text": ["Hello", "World"]}) + y = { + "series": make_series(make_df, [0, 1]), + "list": [0, 1], + "array": np.array([0, 1]), + }[target] + transformer = TextFeatures(variables=["text"], features=["char_count"]) + Xt = transformer.fit(X, y).transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"text": ["Hello", "World"], "text_char_count": [5, 5]} -def test_missing_variable_raises_error(): - X = pd.DataFrame({"text": ["Hello"]}) +def test_error_if_variable_not_in_dataframe(make_df): + X = make_df({"text": ["Hello"]}) transformer = TextFeatures(variables=["nonexistent"]) - with pytest.raises(ValueError, match="not present in the dataframe"): + msg = "Variables {'nonexistent'} are not present in the dataframe." + with pytest.raises(ValueError, match=re.escape(msg)): transformer.fit(X) -@pytest.mark.parametrize("variables", ["Age", "Marks", "dob"]) -def test_no_text_columns_raises_error(df_vartypes, variables): - transformer = TextFeatures(variables=variables) - with pytest.raises(ValueError, match="not object or string"): - transformer.fit(df_vartypes) - - -def test_nan_handling_raise_error_fit(df_na): - transformer = TextFeatures( - variables=["City"], features=["char_count"], missing_values="raise" +@pytest.mark.parametrize( + "variable, values", + [ + ("Age", [20, 21]), + ("Marks", [0.9, 0.8]), + ("dob", [datetime(2020, 2, 24), datetime(2020, 2, 25)]), + ], +) +def test_error_if_variable_not_text(make_df, variable, values): + X = make_df({"Name": ["tom", "nick"], variable: values}) + transformer = TextFeatures(variables=["Name", variable]) + msg = ( + f"Variables ['{variable}'] are not object or string. " + "Please provide text variables only." ) - msg = "`missing_values='ignore'` when initialising this transformer" - with pytest.raises(ValueError, match=msg): - transformer.fit(df_na) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(X) -# ============================================================================== -# TRANSFORM TESTS - GENERAL -# ============================================================================== +def test_categorical_variables_with_string_categories(make_df): + X = make_df({"text": ["Hello World", "Hi", "Hello World"]}) + X = nw.from_native(X).with_columns(nw.col("text").cast(nw.Categorical)) + transformer = TextFeatures(variables=["text"], features=["word_count"]) + Xt = transformer.fit_transform(X.to_native()) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "text": ["Hello World", "Hi", "Hello World"], + "text_word_count": [2, 1, 2], + } -def test_transform_on_new_data(): - X_train = pd.DataFrame({"text": ["Hello World", "Foo Bar"]}) - X_test = pd.DataFrame({"text": ["New Data", "Test 123"]}) - transformer = TextFeatures( - variables=["text"], features=["char_count", "has_digits"] +def test_error_if_categories_are_not_strings(): + # polars categories are always strings + X = pd.DataFrame({"text": pd.Series([1, 2], dtype="category")}) + transformer = TextFeatures(variables=["text"]) + msg = ( + "Variables ['text'] are not object or string. " + "Please provide text variables only." ) - transformer.fit(X_train) - X_tr = transformer.transform(X_test) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(X) - assert X_tr["text_char_count"].tolist() == [7, 7] - assert X_tr["text_has_digits"].tolist() == [0, 1] +def test_error_if_missing_values_in_fit(make_df): + X = make_df({"text": ["Hello", None, "World"]}) + transformer = TextFeatures(variables=["text"], missing_values="raise") + msg = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer or set the parameter " + "`missing_values='ignore'` when initialising this transformer." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.fit(X) -def test_nan_handling_raise_error_transform(): - X_train = pd.DataFrame({"text": ["Hello", "World"]}) - X_test = pd.DataFrame({"text": ["Hello", None, "World"]}) - transformer = TextFeatures( - variables=["text"], features=["char_count"], missing_values="raise" + +def test_error_if_missing_values_in_transform(make_df): + transformer = TextFeatures(variables=["text"], missing_values="raise") + transformer.fit(make_df({"text": ["Hello", "World"]})) + msg = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer or set the parameter " + "`missing_values='ignore'` when initialising this transformer." ) - transformer.fit(X_train) - msg = "`missing_values='ignore'` when initialising this transformer" - with pytest.raises(ValueError, match=msg): - transformer.transform(X_test) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.transform(make_df({"text": ["Hello", None, "World"]})) -def test_nan_handling(): - X = pd.DataFrame({"text": ["Hello", None, "World"]}) +def test_missing_values_are_treated_as_empty_strings(make_df): + X = make_df({"text": ["Hello", None, "World"]}) transformer = TextFeatures(variables=["text"], features=["char_count"]) - X_tr = transformer.fit_transform(X) + Xt = transformer.fit_transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "text": ["Hello", "", "World"], + "text_char_count": [5, 0, 5], + } + assert frame_to_dict(X) == {"text": ["Hello", None, "World"]} + - # NaN should be filled with empty string, resulting in char_count of 0 - assert X_tr["text_char_count"].tolist() == [5, 0, 5] +def test_missing_values_raise_returns_same_values(make_df): + X = make_df({"text": ["Hello World", "Hi"]}) + transformer = TextFeatures( + variables=["text"], features=["word_count"], missing_values="raise" + ) + Xt = transformer.fit_transform(X) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "text": ["Hello World", "Hi"], + "text_word_count": [2, 1], + } -def test_default_all_features(): - """Test extracting all features with default parameters.""" - X = pd.DataFrame({"text": ["Hello World!", "Python 123", "AI"]}) + +def test_error_if_not_fitted(make_df): transformer = TextFeatures(variables=["text"]) - X_tr = transformer.fit_transform(X) + msg = ( + "This TextFeatures instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + transformer.transform(make_df({"text": ["Hello"]})) - # Spot check a few features to ensure they were added and computed - assert X_tr["text_char_count"].tolist() == [11, 9, 2] - assert X_tr["text_word_count"].tolist() == [2, 2, 1] - assert X_tr["text_digit_count"].tolist() == [0, 3, 0] + +def test_error_if_transform_gets_different_number_of_columns(make_df): + transformer = TextFeatures(variables=["text"]).fit(make_df({"text": ["Hello"]})) + msg = ( + "The number of columns in this dataset is different from the one used to " + "fit this transformer (when using the fit() method)." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.transform(make_df({"text": ["Hello"], "other": [1]})) -def test_specific_features(): - """Test extracting specific features only.""" - X = pd.DataFrame({"text": ["Hello", "World"]}) +def test_transform_on_new_data(make_df): transformer = TextFeatures( - variables=["text"], features=["char_count", "word_count"] + variables=["text"], features=["char_count", "has_digits"] ) - X_tr = transformer.fit_transform(X) + transformer.fit(make_df({"text": ["Hello World", "Foo Bar"]})) + Xt = transformer.transform(make_df({"text": ["New Data", "Test 123"]})) - # Check only specified features are extracted - assert X_tr.columns.tolist() == ["text", "text_char_count", "text_word_count"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "text": ["New Data", "Test 123"], + "text_char_count": [7, 7], + "text_has_digits": [0, 1], + } -def test_specific_variables(): - """Test extracting features from specific variables only.""" - X = pd.DataFrame( - {"text1": ["Hello", "World"], "text2": ["Foo", "Bar"], "numeric": [1, 2]} - ) - transformer = TextFeatures(variables=["text1"], features=["char_count"]) - X_tr = transformer.fit_transform(X) +def test_transform_reorders_columns_as_in_fit(make_df): + transformer = TextFeatures(variables=["text"], features=["char_count"]) + transformer.fit(make_df({"text": ["Hello"], "other": [1]})) + Xt = transformer.transform(make_df({"other": [2], "text": ["Hi"]})) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"text": ["Hi"], "other": [2], "text_char_count": [2]} + + +def test_default_extracts_all_features(make_df): + X = make_df({"text": ["Hello World!", "Python 123", "AI"]}) + Xt = TextFeatures(variables=["text"]).fit_transform(X) - # Only text1 should have features extracted - assert X_tr.columns.tolist() == ["text1", "text2", "numeric", "text1_char_count"] + assert isinstance(Xt, make_df) + assert list(Xt.columns) == ["text"] + [f"text_{f}" for f in TEXT_FEATURES] -def test_drop_original(): - """Test drop_original parameter.""" - X = pd.DataFrame({"text": ["Hello", "World"], "other": [1, 2]}) +def test_only_selected_variables_and_features_are_added(make_df): + X = make_df({"a": ["Hello", "World"], "b": ["Foo", "Bar"], "numeric": [1, 2]}) + transformer = TextFeatures( + variables=["b", "a"], features=["word_count", "is_empty"] + ) + Xt = transformer.fit_transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "a": ["Hello", "World"], + "b": ["Foo", "Bar"], + "numeric": [1, 2], + "b_word_count": [1, 1], + "b_is_empty": [0, 0], + "a_word_count": [1, 1], + "a_is_empty": [0, 0], + } + + +def test_drop_original(make_df): + X = make_df({"text": ["Hello", "World"], "other": [1, 2]}) transformer = TextFeatures( variables=["text"], features=["char_count"], drop_original=True ) - X_tr = transformer.fit_transform(X) + Xt = transformer.fit_transform(X) - assert X_tr.columns.tolist() == ["other", "text_char_count"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"other": [1, 2], "text_char_count": [5, 5]} -def test_string_variable_input(): - """Test that passing a single string variable works (auto-converted to list).""" - X = pd.DataFrame({"text": ["Hello", "World"], "other": ["A", "B"]}) - transformer = TextFeatures(variables="text", features=["char_count"]) - X_tr = transformer.fit_transform(X) +@pytest.mark.parametrize("feature", list(TEXT_FEATURES.keys())) +def test_feature_values(make_df, feature): + X = make_df({"text": TEXT}) + Xt = TextFeatures(variables=["text"], features=[feature]).fit_transform(X) - assert transformer.variables_ == ["text"] - assert X_tr.columns.tolist() == ["text", "other", "text_char_count"] - assert X_tr["text_char_count"].tolist() == [5, 5] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt)[f"text_{feature}"] == pytest.approx(EXPECTED[feature]) -def test_multiple_text_columns(): - """Test extracting features from multiple text columns.""" - X = pd.DataFrame({"a": ["Hello", "World"], "b": ["Foo", "Bar"]}) - transformer = TextFeatures( - variables=["a", "b"], features=["char_count", "word_count"] +@pytest.mark.parametrize("feature", list(TEXT_FEATURES.keys())) +def test_feature_values_on_edge_cases(make_df, feature): + X = make_df({"text": TEXT_EDGE_CASES}) + Xt = TextFeatures(variables=["text"], features=[feature]).fit_transform(X) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt)[f"text_{feature}"] == pytest.approx( + EXPECTED_EDGE_CASES[feature] ) - X_tr = transformer.fit_transform(X) - - assert X_tr.columns.tolist() == [ - "a", - "b", - "a_char_count", - "a_word_count", - "b_char_count", - "b_word_count", - ] -# ============================================================================== -# TRANSFORM - TEST TEXT FEATURES -# ============================================================================== +@pytest.mark.parametrize("feature", list(TEXT_FEATURES.keys())) +def test_feature_values_on_other_backends(monkeypatch, feature): + # makes polars take the path used by backends other than pandas and polars + backend_checks = SimpleNamespace( + is_pandas_dataframe=lambda X: False, is_polars_dataframe=lambda X: False + ) + monkeypatch.setattr(text_features, "nwd", backend_checks) + X = pl.DataFrame({"text": TEXT_EDGE_CASES}) + Xt = TextFeatures(variables=["text"], features=[feature]).fit_transform(X) + + assert isinstance(Xt, pl.DataFrame) + assert frame_to_dict(Xt)[f"text_{feature}"] == pytest.approx( + EXPECTED_EDGE_CASES[feature] + ) -@pytest.fixture(scope="module") -def df_text(): - df = pd.DataFrame( +def test_lexical_diversity_is_unique_words_over_total_words(make_df): + X = make_df( { "text": [ - "Hello World!", - "HELLO", - "12345", - "e.g. i.e.", - " ", - " trailing ", - "abc...", - "", - None, - "A? B! C.", - "HeLLo", - "Hi! @#", - "A1b2 C3d4!@#$", - "???", - "i.e., this is wrong", - "Is 1 > 2? No, 100%!", - "Hello. World", - "Hello. World.", - "Hello... World!?!", - "This is a proper sentence containing " - "supercalifragilisticexpialidocious and exceptionally long words.", + "the cat sat on the mat", # 6 words, 5 unique + "good good good good", # 4 words, 1 unique + "all words here are distinct", # 5 words, 5 unique ] } ) - return df - - -def test_whitespace_features(df_text): - text_features = ["whitespace_count", "whitespace_ratio"] - transformer = TextFeatures(variables=["text"], features=text_features) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_whitespace_count"].tolist() == [ - 1, - 0, - 0, - 1, - 3, - 2, - 0, - 0, - 0, - 2, - 0, - 1, - 1, - 0, - 3, - 5, - 1, - 1, - 1, - 10, - ] - assert X_tr["text_whitespace_ratio"].tolist() == [ - 0.08333333333333333, - 0.0, - 0.0, - 0.1111111111111111, - 1.0, - 0.2, - 0.0, - 0.0, - 0.0, - 0.25, - 0.0, - 0.16666666666666666, - 0.07692307692307693, - 0.0, - 0.15789473684210525, - 0.2631578947368421, - 0.08333333333333333, - 0.07692307692307693, - 0.058823529411764705, - 0.09900990099009901, - ] - + transformer = TextFeatures(variables=["text"], features=["lexical_diversity"]) + Xt = transformer.fit_transform(X) -def test_digit_features(df_text): - transformer = TextFeatures( - variables=["text"], features=["digit_count", "digit_ratio", "has_digits"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt)["text_lexical_diversity"] == pytest.approx( + [5 / 6, 1 / 4, 1] ) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_digit_count"].tolist() == [ - 0, - 0, - 5, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 4, - 0, - 0, - 5, - 0, - 0, - 0, - 0, - ] - assert X_tr["text_digit_ratio"].tolist() == [ - 0.0, - 0.0, - 1.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.3333333333333333, - 0.0, - 0.0, - 0.35714285714285715, - 0.0, - 0.0, - 0.0, - 0.0, - ] - assert X_tr["text_has_digits"].tolist() == [ - 0, - 0, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 0, - ] -def test_uppercase_features(df_text): - transformer = TextFeatures( - variables=["text"], - features=[ - "uppercase_count", - "uppercase_ratio", - "has_uppercase", - "starts_with_uppercase", - ], - ) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_uppercase_count"].tolist() == [ - 2, - 5, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 3, - 3, - 1, - 2, - 0, - 0, - 2, - 2, - 2, - 2, - 1, - ] - assert X_tr["text_uppercase_ratio"].tolist() == [ - 0.18181818181818182, - 1.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.5, - 0.6, - 0.2, - 0.16666666666666666, - 0.0, - 0.0, - 0.14285714285714285, - 0.18181818181818182, - 0.16666666666666666, - 0.125, - 0.01098901098901099, - ] - assert X_tr["text_has_uppercase"].tolist() == [ - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 1, - 1, - 0, - 0, - 1, - 1, - 1, - 1, - 1, - ] - assert X_tr["text_starts_with_uppercase"].tolist() == [ - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 1, - 1, - 0, - 0, - 1, - 1, - 1, - 1, - 1, - ] - +def test_output_dtypes(make_df): + X = make_df({"text": ["Hello World", "Hi"]}) + features = ["char_count", "whitespace_ratio", "has_digits", "unique_word_count"] + Xt = TextFeatures(variables=["text"], features=features).fit_transform(X) -def test_punctuation_features(df_text): - transformer = TextFeatures( - variables=["text"], features=["special_char_count", "ends_with_punctuation"] - ) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_special_char_count"].tolist() == [ - 1, - 0, - 0, - 4, - 0, - 0, - 3, - 0, - 0, - 3, - 0, - 3, - 4, - 3, - 3, - 5, - 1, - 2, - 6, - 1, - ] - assert X_tr["text_ends_with_punctuation"].tolist() == [ - 1, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 1, - 0, - 0, - 0, - 1, - 0, - 1, - 0, - 1, - 1, - 1, + schema = nw.from_native(Xt).schema + assert [schema[f"text_{f}"] for f in features] == [ + nw.Int64, + nw.Float64, + nw.Int64, + nw.Int64, ] -def test_word_features(df_text): +@pytest.mark.parametrize( + "drop_original, expected", + [ + (False, ["text", "other", "text_char_count", "text_word_count"]), + (True, ["other", "text_char_count", "text_word_count"]), + ], +) +def test_get_feature_names_out(make_df, drop_original, expected): + X = make_df({"text": ["Hello"], "other": [1]}) transformer = TextFeatures( variables=["text"], - features=[ - "word_count", - "unique_word_count", - "lexical_diversity", - "avg_word_length", - ], + features=["char_count", "word_count"], + drop_original=drop_original, ) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_word_count"].tolist() == [ - 2, - 1, - 1, - 2, - 0, - 1, - 1, - 0, - 0, - 3, - 1, - 2, - 2, - 1, - 4, - 6, - 2, - 2, - 2, - 11, - ] - assert X_tr["text_unique_word_count"].tolist() == [ - 2, - 1, - 1, - 2, - 0, - 1, - 1, - 0, - 0, - 3, - 1, - 2, - 2, - 1, - 4, - 6, - 2, - 2, - 2, - 11, - ] - assert X_tr["text_lexical_diversity"].tolist() == [ - 1.0, - 1.0, - 1.0, - 1.0, - 0.0, - 1.0, - 1.0, - 0.0, - 0.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - ] - assert X_tr["text_avg_word_length"].tolist() == [ - 6.0, - 5.0, - 5.0, - 4.5, - 0.0, - 8.0, - 6.0, - 0.0, - 0.0, - 2.6666666666666665, - 5.0, - 3.0, - 6.5, - 3.0, - 4.75, - 3.1666666666666665, - 6.0, - 6.5, - 8.5, - 9.181818181818182, - ] + Xt = transformer.fit_transform(X) + assert transformer.get_feature_names_out() == expected + assert list(Xt.columns) == expected -def test_basic_features(df_text): - transformer = TextFeatures( - variables=["text"], - features=[ - "char_count", - "sentence_count", - "letter_count", - "lowercase_count", - "is_empty", - ], - ) - X_tr = transformer.fit_transform(df_text) - assert X_tr["text_char_count"].tolist() == [ - 11, - 5, - 5, - 8, - 0, - 8, - 6, - 0, - 0, - 6, - 5, - 5, - 12, - 3, - 16, - 14, - 11, - 12, - 16, - 91, - ] - assert X_tr["text_sentence_count"].tolist() == [ - 1, - 0, - 0, - 4, - 0, - 0, - 1, - 0, - 0, - 3, - 0, - 1, - 1, - 1, - 2, - 2, - 1, - 2, - 2, - 1, - ] - assert X_tr["text_letter_count"].tolist() == [ - 10, - 5, - 0, - 4, - 0, - 8, - 3, - 0, - 0, - 3, - 5, - 2, - 4, - 0, - 13, - 4, - 10, - 10, - 10, - 90, - ] - assert X_tr["text_lowercase_count"].tolist() == [ - 8, - 0, - 0, - 4, - 0, - 8, - 3, - 0, - 0, - 0, - 2, - 1, - 2, - 0, - 13, - 2, - 8, - 8, - 8, - 89, - ] - assert X_tr["text_is_empty"].tolist() == [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 1, - 1, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, + +@pytest.mark.parametrize( + "input_features", [["text", "other"], np.array(["text", "other"])] +) +def test_get_feature_names_out_with_input_features(make_df, input_features): + X = make_df({"text": ["Hello"], "other": [1]}) + transformer = TextFeatures(variables=["text"], features=["char_count"]).fit(X) + assert transformer.get_feature_names_out(input_features) == [ + "text", + "other", + "text_char_count", ] -# ============================================================================== -# OTHER METHOD TESTS -# ============================================================================== +@pytest.mark.parametrize("input_features", [["other", "text"], ["text"]]) +def test_error_if_input_features_not_feature_names_in(make_df, input_features): + X = make_df({"text": ["Hello"], "other": [1]}) + transformer = TextFeatures(variables=["text"], features=["char_count"]).fit(X) + msg = "input_features is not equal to feature_names_in_" + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.get_feature_names_out(input_features) -def test_get_feature_names_out(): - X = pd.DataFrame({"text": ["Hello"], "other": [1]}) - transformer = TextFeatures( - variables=["text"], features=["char_count", "word_count"] - ) - transformer.fit(X) +@pytest.mark.parametrize("input_features", ["text", 1, {"text": 1}]) +def test_error_if_input_features_not_list_or_array(make_df, input_features): + X = make_df({"text": ["Hello"], "other": [1]}) + transformer = TextFeatures(variables=["text"], features=["char_count"]).fit(X) + msg = f"input_features must be a list or an array. Got {input_features} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + transformer.get_feature_names_out(input_features) - feature_names = transformer.get_feature_names_out() - expected_features = ["text", "other", "text_char_count", "text_word_count"] - assert feature_names == expected_features +def test_integer_column_names(): + X = pd.DataFrame({0: [1, 2], "text": ["Hello World", None], 1: ["a", "b"]}) + transformer = TextFeatures(variables=["text"], features=["word_count"]) + Xt = transformer.fit_transform(X) -def test_get_feature_names_out_with_drop(): - """Test get_feature_names_out with drop_original=True.""" - X = pd.DataFrame({"text": ["Hello"], "other": [1]}) - transformer = TextFeatures( - variables=["text"], features=["char_count"], drop_original=True + expected = pd.DataFrame( + { + 0: [1, 2], + "text": ["Hello World", ""], + 1: ["a", "b"], + "text_word_count": [2, 0], + } ) - transformer.fit(X) + pd.testing.assert_frame_equal(Xt, expected) + assert transformer.get_feature_names_out() == [0, "text", 1, "text_word_count"] - feature_names = transformer.get_feature_names_out() - expected_features = ["other", "text_char_count"] - assert feature_names == expected_features +def test_pandas_index_is_kept(): + X = pd.DataFrame({"text": ["Hello World", "Hi", "Hey"]}, index=[10, 10, 3]) + transformer = TextFeatures( + variables=["text"], features=["char_count", "unique_word_count"] + ) + Xt = transformer.fit_transform(X) -def test_lexical_diversity_is_unique_words_over_total_words(): - X = pd.DataFrame( + expected = pd.DataFrame( { - "text": [ - "the cat sat on the mat", # 6 words, 5 unique - "good good good good", # 4 words, 1 unique - "all words here are distinct", # 5 words, 5 unique - ] - } + "text": ["Hello World", "Hi", "Hey"], + "text_char_count": [10, 2, 3], + "text_unique_word_count": [2, 1, 1], + }, + index=[10, 10, 3], ) - transformer = TextFeatures(variables=["text"], features=["lexical_diversity"]) - X_tr = transformer.fit_transform(X) - - assert X_tr["text_lexical_diversity"].tolist() == [5 / 6, 1 / 4, 1.0] - # a ratio of unique words to total words never exceeds 1 - assert (X_tr["text_lexical_diversity"] <= 1.0).all() + pd.testing.assert_frame_equal(Xt, expected)