From e3ccfb40333887e123920c2512e5736ee42772c4 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:32:14 +0200 Subject: [PATCH 1/2] Migrate TextFeatures to narwhals, add polars support TextFeatures now accepts pandas, polars and other dataframes supported by narwhals, and returns the same type it receives. All features are defined once in terms of a few text statistics, computed with pandas string methods and Python loops for pandas, polars string methods for polars, and narwhals expressions for other backends. Pandas outputs are identical to before and the transform is about 3x faster. Co-Authored-By: Claude Opus 5 --- docs/user_guide/text/TextFeatures.rst | 77 +- feature_engine/text/text_features.py | 376 +++++--- tests/test_text/test_text_features.py | 1135 ++++++++++--------------- 3 files changed, 773 insertions(+), 815 deletions(-) diff --git a/docs/user_guide/text/TextFeatures.rst b/docs/user_guide/text/TextFeatures.rst index cf31dd0b9..36e5c02ed 100644 --- a/docs/user_guide/text/TextFeatures.rst +++ b/docs/user_guide/text/TextFeatures.rst @@ -38,27 +38,31 @@ 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**: Number of characters from the first to the last non-whitespace + character, including the spaces between words, divided by the number of words - **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 +164,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 @@ -279,6 +283,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..df8f3a64d 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,175 @@ 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.strip_length() / 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 strip_length(self): + return self.text.str.strip().str.len() + + 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 strip_length(self): + return self.text.str.strip_chars(_WHITESPACE).str.len_chars() + + 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 +208,26 @@ 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': Number of characters, from the first to the last + non-whitespace character, divided by the number of words - '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 +235,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. @@ -160,7 +307,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 +314,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 +334,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,75 +373,96 @@ 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" ) + + # 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[self.variables_] = X[self.variables_].fillna("") + X_new = self._transform_expressions(nw_X, nw, _NarwhalsText).to_native() + + return X_new - # reorder variables to match train set - X = X[self.feature_names_in_] + 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_}) - # Extract features for each text variable + new_features = [] 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]) + statistics = _PandasText(X_new[var], native_namespace) + new_features += [ + TEXT_FEATURES[feature](statistics).rename(f"{var}_{feature}") + for feature in self.features_ + ] + + 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_) - # Fill any NaN values resulting from computation with 0 - X[new_col_name] = X[new_col_name].fillna(0) + return X_new + + 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_ + ] - if self.drop_original: - X = X.drop(columns=self.variables_) + 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 + return X_new def get_feature_names_out(self, input_features=None) -> List[str]: """ @@ -321,7 +481,7 @@ def get_feature_names_out(self, input_features=None) -> List[str]: check_is_fitted(self) # Start with original features - if self.drop_original: + if self.drop_original is True: feature_names = [ f for f in self.feature_names_in_ if f not in self.variables_ ] diff --git a/tests/test_text/test_text_features.py b/tests/test_text/test_text_features.py index 827ff6391..6cdf29b9e 100644 --- a/tests/test_text/test_text_features.py +++ b/tests/test_text/test_text_features.py @@ -1,798 +1,541 @@ +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": [ + 6, 5, 5, 9 / 2, 0, 8, 6, 0, 0, 8 / 3, 5, 3, 13 / 2, 3, 19 / 4, 19 / 6, 6, + 13 / 2, 17 / 2, 101 / 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, 6, 5, 5, 4, 5 / 3, 2, 2, 2, 11 / 3, 9 / 2, 19 / 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, - ] - - -def test_digit_features(df_text): - transformer = TextFeatures( - variables=["text"], features=["digit_count", "digit_ratio", "has_digits"] - ) - 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, - ] - + transformer = TextFeatures(variables=["text"], features=["lexical_diversity"]) + Xt = transformer.fit_transform(X) -def test_uppercase_features(df_text): - transformer = TextFeatures( - variables=["text"], - features=[ - "uppercase_count", - "uppercase_ratio", - "has_uppercase", - "starts_with_uppercase", - ], + 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_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_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, - ] - +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_word_features(df_text): - transformer = TextFeatures( - variables=["text"], - features=[ - "word_count", - "unique_word_count", - "lexical_diversity", - "avg_word_length", - ], - ) - 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, + 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_basic_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=[ - "char_count", - "sentence_count", - "letter_count", - "lowercase_count", - "is_empty", - ], + features=["char_count", "word_count"], + drop_original=drop_original, ) - 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, - ] + Xt = transformer.fit_transform(X) + assert transformer.get_feature_names_out() == expected + assert list(Xt.columns) == expected -# ============================================================================== -# OTHER METHOD TESTS -# ============================================================================== +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(): - X = pd.DataFrame({"text": ["Hello"], "other": [1]}) - transformer = TextFeatures( - variables=["text"], features=["char_count", "word_count"] + expected = pd.DataFrame( + { + 0: [1, 2], + "text": ["Hello World", ""], + 1: ["a", "b"], + "text_word_count": [2, 0], + } ) - transformer.fit(X) - - feature_names = transformer.get_feature_names_out() - expected_features = ["text", "other", "text_char_count", "text_word_count"] - assert feature_names == expected_features + pd.testing.assert_frame_equal(Xt, expected) + assert transformer.get_feature_names_out() == [0, "text", 1, "text_word_count"] -def test_get_feature_names_out_with_drop(): - """Test get_feature_names_out with drop_original=True.""" - X = pd.DataFrame({"text": ["Hello"], "other": [1]}) +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"], drop_original=True + variables=["text"], features=["char_count", "unique_word_count"] ) - transformer.fit(X) - - feature_names = transformer.get_feature_names_out() - expected_features = ["other", "text_char_count"] - assert feature_names == expected_features + 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) From 4ddcf2ade31a2ec62456970a43b9baf6524044dc Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 13:13:21 +0200 Subject: [PATCH 2/2] Exclude spaces from avg_word_length and use the shared get_feature_names_out in TextFeatures Co-Authored-By: Claude Opus 5 --- docs/user_guide/text/TextFeatures.rst | 11 +++--- feature_engine/text/text_features.py | 49 +++++---------------------- tests/test_text/test_text_features.py | 39 +++++++++++++++++++-- 3 files changed, 49 insertions(+), 50 deletions(-) diff --git a/docs/user_guide/text/TextFeatures.rst b/docs/user_guide/text/TextFeatures.rst index 36e5c02ed..1ad66bd4b 100644 --- a/docs/user_guide/text/TextFeatures.rst +++ b/docs/user_guide/text/TextFeatures.rst @@ -41,8 +41,7 @@ Text features - **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**: Number of characters from the first to the last non-whitespace - character, including the spaces between words, divided by the number of words +- **avg_word_length**: Average number of characters per word - **digit_count**: Number of digit characters - **letter_count**: Number of letters a-z and A-Z - **uppercase_count**: Number of uppercase letters A-Z @@ -225,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 diff --git a/feature_engine/text/text_features.py b/feature_engine/text/text_features.py index df8f3a64d..3d492edf6 100644 --- a/feature_engine/text/text_features.py +++ b/feature_engine/text/text_features.py @@ -35,7 +35,8 @@ "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.strip_length() / t.word_count.clip(1), + "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]"), @@ -83,9 +84,6 @@ 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 strip_length(self): - return self.text.str.strip().str.len() - def count(self, pattern): if ("count", pattern) not in self._counts: self._counts[("count", pattern)] = self.text.str.count(pattern) @@ -144,9 +142,6 @@ def unique_word_count(self): .cast(self._ns.Int64) ) - def strip_length(self): - return self.text.str.strip_chars(_WHITESPACE).str.len_chars() - 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 @@ -211,8 +206,7 @@ class TextFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): - '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': Number of characters, from the first to the last - non-whitespace character, divided by the number of words + - 'avg_word_length': Average number of characters per word - 'digit_count': Number of digit characters - 'letter_count': Number of letters a-z and A-Z - 'uppercase_count': Number of uppercase letters A-Z @@ -288,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) @@ -464,33 +456,8 @@ def _transform_expressions(self, X, namespace, text_class): return X_new - def get_feature_names_out(self, input_features=None) -> List[str]: - """ - Get output feature names for transformation. - - Parameters - ---------- - input_features : array-like of str or None, default=None - Input features. If ``None``, uses ``feature_names_in_``. - - Returns - ------- - feature_names_out : list of str - Output feature names. - """ - check_is_fitted(self) - - # Start with original features - if self.drop_original is True: - feature_names = [ - f for f in self.feature_names_in_ if f not in self.variables_ - ] - 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}") - - 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 6cdf29b9e..26568748f 100644 --- a/tests/test_text/test_text_features.py +++ b/tests/test_text/test_text_features.py @@ -65,8 +65,8 @@ "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": [ - 6, 5, 5, 9 / 2, 0, 8, 6, 0, 0, 8 / 3, 5, 3, 13 / 2, 3, 19 / 4, 19 / 6, 6, - 13 / 2, 17 / 2, 101 / 11, + 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], @@ -102,7 +102,9 @@ "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, 6, 5, 5, 4, 5 / 3, 2, 2, 2, 11 / 3, 9 / 2, 19 / 6], + "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], @@ -506,6 +508,37 @@ def test_get_feature_names_out(make_df, drop_original, expected): assert list(Xt.columns) == expected +@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", + ] + + +@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) + + +@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) + + 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"])