Miscellaneous

Serializable

Source code in src\survey_kit\serializable.py

class Serializable:
    _save_exclude_items = []
    _save_suffix = "serial"
    _registry = {}

    try:
        import numpy as np

        _BASE_TYPES = (int, float, str, bool, np.float64)
    except ImportError:
        _BASE_TYPES = (int, float, str, bool)

    def __init__(self):
        self.__fully_serializable__ = True

    def __init_subclass__(cls, **kwargs):
        """Automatically register subclasses when they're defined"""
        super().__init_subclass__(**kwargs)
        if hasattr(cls, "_save_suffix") and cls._save_suffix:
            Serializable._registry[cls._save_suffix] = cls
            logger.debug(f"Registered {cls.__name__} with suffix '{cls._save_suffix}'")

    @classmethod
    def _init_from_dict(
        cls, data: dict, init_kwargs: dict | None = None, **additional_kwargs
    ):
        init_keys = list(inspect.signature(cls.__init__).parameters.keys())
        init_keys.remove("self")
        if init_kwargs is None:
            init_kwargs = {}
        remaining_kwargs = {}
        for key, value in data.items():
            if key in init_keys:
                init_kwargs[key] = value
            else:
                remaining_kwargs[key] = value

        obj = cls(**init_kwargs, **additional_kwargs)

        for keyi, valuei in remaining_kwargs.items():
            if (keyi not in obj._reserved_keys()) and (
                keyi not in obj._save_exclude_items
            ):
                setattr(obj, keyi, valuei)

        return obj

    def __hash__(self):
        d_save = self.to_dict()
        if self.__fully_serializable__:
            save_caller = json
        else:
            save_caller = pickle

        return hash(save_caller.dumps(d_save))

    def save(self, path: str, quietly: bool = True) -> None:
        """
        Save a serializable object to

        Parameters
        ----------
        path : str, optional
            Path of object.  Can exclude the suffix.
        quietly: bool, optional
            No message to console/log
        """

        self.__serializable_n_dfs = 0

        d_save = self.to_dict()

        if type(path) is not str:
            path = str(path)
        if not path.endswith(f".{self._save_suffix}"):
            folder_path = f"{path}.{self._save_suffix}"
        else:
            folder_path = path

        #   logger.info(folder_path)

        if os.path.isdir(folder_path):
            logger.info("Removing existing directory " + folder_path)
            shutil.rmtree(folder_path)

        #   Make the path to save everything
        create_folders_if_needed([folder_path])
        #   os.makedirs(folder_path)

        #   Save the data
        dfs = d_save["__serialized_dfs__"]
        if len(dfs):
            self._save_dfs(folder_path=folder_path, dfs=dfs, quietly=quietly)

        if self.__fully_serializable__:
            suffix = "json"
            save_caller = json
            w = "w"
        else:
            suffix = "pkl"
            save_caller = pickle
            w = "wb"

        path_object = f"{folder_path}/{self._save_suffix}.{suffix}"
        with open(path_object, w) as f:
            save_caller.dump(d_save, f)

    def _save_dfs(self, folder_path: str, dfs: dict, quietly: bool = True) -> None:
        for valuei in dfs.values():
            dfi = valuei["df"]
            pathi = valuei["path"]

            path_save = os.path.normpath(f"{folder_path}/{pathi}")
            create_folders_if_needed([os.path.dirname(path_save)])

            dfi_nw = nw.from_native(dfi)
            d_metadata = dict(engine=nw.get_native_namespace(dfi_nw).__package__)
            SerializableDictionary(d_metadata).save(path_save)
            if isinstance(dfi_nw, nw.LazyFrame):
                dfi_nw.sink_parquet(path_save)
            else:
                dfi_nw.write_parquet(path_save)

            del valuei["df"]

    @classmethod
    def delete(cls, path: str):
        cls.load(path=path, delete=True, delete_only=True)

    @classmethod
    def load(
        cls,
        path: str = "",
        delete: bool = False,
        delete_only: bool = False,
        init_kwargs: dict | None = None,
        **df_kwargs,
    ) -> Serializable | None:
        """
        Load a serializable object from disk

        Parameters
        ----------
        path : str, optional
            Path of object.  Can exclude the suffix.
        delete : bool, optional
            Delete after load? The default is False.
        delete_only : bool, optional
            Don't load, just delete. The default is False.
        init_kwargs : dict, optional
            Arguments to pass to init function
        df_kwargs : dict, optional
            Arguments to pass to narwhals scan_parquet for dataframe loading
        Returns
        -------
        Serializable
        """

        if type(path) is not str:
            path = str(path)

        # If called on Serializable base class, auto-detect the subclass
        if cls == Serializable:
            # Find which suffix matches
            detected_class = None
            for suffix, registered_cls in Serializable._registry.items():
                if path.endswith(f".{suffix}") or os.path.isdir(f"{path}.{suffix}"):
                    detected_class = registered_cls
                    break

            if detected_class is None:
                message = f"Could not detect serializable type for path: {path}"
                logger.error(message)
                raise Exception(message)

            # Delegate to the correct class
            return detected_class.load(
                path=path, delete=delete, delete_only=delete_only, **df_kwargs
            )

        if not path.endswith(f".{cls._save_suffix}"):
            folder_path = f"{path}.{cls._save_suffix}"
        else:
            folder_path = path

        if not delete_only:
            path_dict = os.path.normpath(f"{folder_path}/{cls._save_suffix}.json")
            path_found = os.path.isfile(path_dict)

            if path_found:
                load_caller = json
            else:
                path_dict = os.path.normpath(f"{folder_path}/{cls._save_suffix}.pkl")
                load_caller = pickle
                path_found = os.path.isfile(path_dict)

            if not path_found:
                message = f"{folder_path} isn't a valid Serializable object"
                logger.error(message)
                raise Exception(message)

            with open(path_dict, "rb") as f:
                d_loaded = load_caller.load(f)

            if len(d_loaded["__serialized_dfs__"]):
                cls._load_dfs(
                    folder_path=folder_path,
                    dfs=d_loaded["__serialized_dfs__"],
                    delete=delete,
                    **df_kwargs,
                )

        if delete or delete_only:
            if os.path.isdir(folder_path):
                logger.info("Removing existing directory " + folder_path)
                shutil.rmtree(folder_path)

        if delete_only:
            return None

        obj = cls.from_dict(d_loaded, init_kwargs=init_kwargs)
        return obj

    @classmethod
    def load_any(cls, path: str = "", delete: bool = False, delete_only: bool = False):
        """
        Pass the root path of a serializable object
            and this will figure out what it is (from the suffix)
            and call that object's loader function.

            For cases where multiple objects work, and I don't want to
            have to do if/else logic by suffix

        Parameters
        ----------
        path : str, optional
            Path of object.  Can exclude the suffix.
        delete : bool, optional
            Delete after load? The default is False.
        delete_only : bool, optional
            Don't load, just delete. The default is False.

        Returns
        -------
        Serializable object of any type

        """

        items = []

        for classi in items:
            suffixi = classi._save_suffix
            if path.endswith(suffixi) or os.path.isdir(f"{path}.{suffixi}"):
                return classi.load(path=path, delete=delete, delete_only=delete_only)

    @classmethod
    def _load_dfs(
        cls, folder_path: str, dfs: dict, delete: bool = False, **df_kwargs
    ) -> None:
        #   Avoid circular import
        from .utilities.dataframe import NarwhalsType

        if df_kwargs is None:
            df_kwargs = {}
        for keyi, valuei in dfs.items():
            pathi = valuei["path"]

            path_load = os.path.normpath(f"{folder_path}/{pathi}")
            d_metadata = SerializableDictionary.load(path_load)
            df_kwargsi = df_kwargs.copy()
            if "backend" not in df_kwargsi:
                df_kwargsi["backend"] = d_metadata["engine"]
                backend = d_metadata["engine"]
            else:
                backend = df_kwargsi["backend"]
            dfi = nw.scan_parquet(path_load, **df_kwargsi)
            if delete:
                dfi = dfi.lazy().collect().lazy_backend(NarwhalsType(backend=backend))

            dfs[keyi]["df"] = dfi.to_native()

    def to_dict(self, dfs: dict | None = None, key_path: str = "") -> dict[str, object]:
        candidate_items = vars(self)

        self.__fully_serializable__ = True
        items = {}
        if dfs is None:
            add_dfs = True
            dfs = {}
        else:
            add_dfs = False

        for keyi, valuei in candidate_items.items():
            self._to_dict_item(
                items=items, dfs=dfs, key_path=key_path, key=keyi, value=valuei
            )

        items["__fully_serializable__"] = self.__fully_serializable__
        items["__class__"] = self.__class__.__qualname__
        if add_dfs:
            items["__serialized_dfs__"] = dfs
        else:
            items["__serialized_dfs__"] = {}

        namespace = vars(inspect.getmodule(self))["__name__"]
        items["__namespace__"] = namespace

        #   logger.info(items)
        return items

    @classmethod
    def from_dict(
        cls,
        data: dict,
        init_kwargs: dict | None = None,
    ) -> object:
        data = cls._from_dict_unpack_item(data, dfs=None, init_kwargs=init_kwargs)

        return data

    @classmethod
    def _from_dict_unpack_item(
        cls,
        item: object,
        dfs: dict | None,
        init_kwargs: dict | None = None,
    ) -> object:
        if dfs is None:
            if "__serialized_dfs__" in item.keys():
                dfs = item["__serialized_dfs__"]
                del item["__serialized_dfs__"]
            else:
                dfs = {}

        typei = type(item)

        if typei is dict:
            if "__enumeration__" in item.keys():
                item = cls._from_dict_unpack_enum(item, dfs)
            elif "__pickled__" in item.keys():
                item = cls._from_dict_unpack_pickled(item, dfs)
            # elif "__polars_expression__" in item.keys():
            #     item = cls._from_dict_unpack_polars_expression(item,
            #                                                    dfs)
            elif cls._is_serializable(item):
                #   Unpack as a separate item
                item = cls._from_dict_unpack_dict(item, dfs, init_kwargs=init_kwargs)

                #   Then turn it into a class
                _namespace = importlib.import_module(item["__namespace__"])
                _class_name = item["__class__"]
                _parent = _namespace
                for classi in _class_name.split("."):
                    # #   Fix for any refactor class name changes
                    # d_refactor = {"MultipleImputationStats":"MultipleImputation"}

                    # classi = d_refactor.get(classi,
                    #                         classi)

                    _class = getattr(_parent, classi)
                    _parent = _class

                #   _class = getattr(_namespace, item['__class__'])

                init_sig = inspect.signature(_class._init_from_dict)
                if "init_kwargs" in init_sig.parameters.keys():
                    item = _class._init_from_dict(item, init_kwargs=init_kwargs)
                else:
                    if init_kwargs is None:
                        init_kwargs = {}
                    item = _class._init_from_dict(item, **init_kwargs)

            else:
                item = cls._from_dict_unpack_dict(item, dfs)
        elif typei is list:
            item = cls._from_dict_unpack_list(item, dfs)
        elif typei is str and item.startswith("__serialized_df_"):
            item = cls._from_dict_assign_df(item, dfs)
        elif typei in [int, float, str, bool]:
            #   Do nothing - don't unpack
            pass

        return item

    @classmethod
    def _from_dict_unpack_dict(
        cls, item: dict, dfs: dict, init_kwargs: dict | None = None
    ) -> object:
        for key, value in item.items():
            if key not in cls._reserved_keys():
                item[key] = cls._from_dict_unpack_item(value, dfs, init_kwargs)

        return item

    @classmethod
    def _from_dict_unpack_list(cls, item: list, dfs: dict):
        for i in range(0, len(item)):
            item[i] = cls._from_dict_unpack_item(item[i], dfs)

        return item

    @classmethod
    def _from_dict_unpack_enum(cls, item: dict, dfs: dict):
        #   Then turn it into a class
        _namespace = importlib.import_module(item["__enum_namespace__"])
        _class_name = item["__enum_class__"]
        _parent = _namespace
        for classi in _class_name.split("."):
            _class = getattr(_parent, classi)
            _parent = _class
        value = item["__enumeration__"]

        return _class(value)

    @classmethod
    def _from_dict_unpack_polars_expression(cls, item: dict, dfs: dict):
        try:
            import polars as pl
        except ImportError:
            raise ImportError("Polars is required for unpacking a polars Expression. ")

        return pl.Expr.deserialize(io.StringIO(item["__polars_expression__"]))

    @classmethod
    def _from_dict_unpack_pickled(cls, item: dict, dfs: dict):
        pickled_data = bytes(item["__pickled__"])
        buffer = io.BytesIO(pickled_data)
        try:
            return pickle.load(buffer)
        except:
            logger.info("Pickled item failed to load")

    @classmethod
    def _from_dict_assign_df(cls, item: str, dfs: dict):
        return dfs[item]["df"]

    def _to_dict_item(
        self,
        items: dict[str, object],
        dfs: dict[str, dict],
        key_path: str,
        key: str,
        value: object,
    ) -> None:
        # logger.info(key)
        # logger.info(value)
        # logger.info(items)
        # logger.info(key_path)

        if (key not in self._reserved_keys()) and (key not in self._save_exclude_items):
            typei = type(value)

            if typei is list:
                d_item = []

                for i in range(0, len(value)):
                    self._to_dict_item(
                        items=d_item,
                        dfs=dfs,
                        key_path=f"{key_path}/{key}/{i}",
                        key=i,
                        value=value[i],
                    )

                self._add_to_items(key=key, value=d_item, items=items)
            elif typei is dict:
                d_item = {}

                for ki, vi in value.items():
                    self._to_dict_item(
                        items=d_item,
                        dfs=dfs,
                        key_path=f"{key_path}/{key}",
                        key=ki,
                        value=vi,
                    )

                self._add_to_items(key=key, value=d_item, items=items)
            elif is_narwhals_compatible(value):
                n_dfs = len(dfs.keys())

                df_name = f"__serialized_df_{n_dfs}"
                dfs[df_name] = {
                    "df": value,
                    "path": f"{key_path}/{key}/{df_name}.parquet",
                }

                self._add_to_items(key=key, value=df_name, items=items)
            elif typei in self._BASE_TYPES:
                self._add_to_items(key=key, value=value, items=items)
            # elif isinstance(value,pl.Expr):
            #     self._add_to_items(key=key,
            #                        value=self._add_to_items_polars_expression(value),
            #                        items=items)
            elif isinstance(value, Serializable):
                value_add = value.to_dict(dfs=dfs, key_path=key_path)
                self.__fully_serializable__ = (
                    self.__fully_serializable__ and value.__fully_serializable__
                )
                self._add_to_items(key=key, value=value_add, items=items)
            elif isinstance(value, Enum):
                value_add = self._add_to_items_enumeration(value)
                self._add_to_items(key=key, value=value_add, items=items)
            elif value is None:
                self._add_to_items(key=key, value=value, items=items)
            else:
                # logger.info(f"{key} not serializable")
                # logger.info(value)
                #   self.__fully_serializable__ = False
                self._add_to_items(
                    key=key, value=self._add_to_items_pickled(value), items=items
                )

    def _add_to_items_enumeration(self, value: Enum) -> dict:
        namespace = value.__module__
        # vars(inspect.getmodule(self))["__name__"]

        # logger.info({
        #     "__enumeration__": value.value,
        #     "__enum_class__": value.__class__.__qualname__,
        #     "__enum_namespace__": namespace,
        # })
        # logger.info("")
        return {
            "__enumeration__": value.value,
            "__enum_class__": value.__class__.__qualname__,
            "__enum_namespace__": namespace,
        }

    def _add_to_items_polars_expression(self, value: pl.Expr) -> dict:
        return {"__polars_expression__": value.meta.serialize()}

    def _add_to_items_pickled(self, value: pl.Expr) -> dict:
        buffer = io.BytesIO()
        pickle.dump(value, buffer)
        value_save = list(buffer.getvalue())
        return {"__pickled__": value_save}

    @staticmethod
    def _add_to_items(key: str, value: object, items: dict | list):
        if type(items) is dict:
            items[key] = value
        elif type(items) is list:
            items.append(value)

    @staticmethod
    def _reserved_keys() -> list[str]:
        return [
            "__class__",
            "__namespace__",
            "__fully_serializable__",
            "__dfs__",
            "__dfpath__",
        ]

    @staticmethod
    def _is_serializable(data: dict):
        return (
            ("__fully_serializable__" in data.keys())
            and ("__class__" in data.keys())
            and ("__namespace__" in data.keys())
        )

load classmethod

load(
    path: str = "",
    delete: bool = False,
    delete_only: bool = False,
    init_kwargs: dict | None = None,
    **df_kwargs,
) -> Serializable | None

Load a serializable object from disk

Parameters:

Name	Type	Description	Default
path	str	Path of object. Can exclude the suffix.	''
delete	bool	Delete after load? The default is False.	False
delete_only	bool	Don't load, just delete. The default is False.	False
init_kwargs	dict	Arguments to pass to init function	None
df_kwargs	dict	Arguments to pass to narwhals scan_parquet for dataframe loading	{}

Returns:

Type	Description
Serializable

Source code in src\survey_kit\serializable.py

@classmethod
def load(
    cls,
    path: str = "",
    delete: bool = False,
    delete_only: bool = False,
    init_kwargs: dict | None = None,
    **df_kwargs,
) -> Serializable | None:
    """
    Load a serializable object from disk

    Parameters
    ----------
    path : str, optional
        Path of object.  Can exclude the suffix.
    delete : bool, optional
        Delete after load? The default is False.
    delete_only : bool, optional
        Don't load, just delete. The default is False.
    init_kwargs : dict, optional
        Arguments to pass to init function
    df_kwargs : dict, optional
        Arguments to pass to narwhals scan_parquet for dataframe loading
    Returns
    -------
    Serializable
    """

    if type(path) is not str:
        path = str(path)

    # If called on Serializable base class, auto-detect the subclass
    if cls == Serializable:
        # Find which suffix matches
        detected_class = None
        for suffix, registered_cls in Serializable._registry.items():
            if path.endswith(f".{suffix}") or os.path.isdir(f"{path}.{suffix}"):
                detected_class = registered_cls
                break

        if detected_class is None:
            message = f"Could not detect serializable type for path: {path}"
            logger.error(message)
            raise Exception(message)

        # Delegate to the correct class
        return detected_class.load(
            path=path, delete=delete, delete_only=delete_only, **df_kwargs
        )

    if not path.endswith(f".{cls._save_suffix}"):
        folder_path = f"{path}.{cls._save_suffix}"
    else:
        folder_path = path

    if not delete_only:
        path_dict = os.path.normpath(f"{folder_path}/{cls._save_suffix}.json")
        path_found = os.path.isfile(path_dict)

        if path_found:
            load_caller = json
        else:
            path_dict = os.path.normpath(f"{folder_path}/{cls._save_suffix}.pkl")
            load_caller = pickle
            path_found = os.path.isfile(path_dict)

        if not path_found:
            message = f"{folder_path} isn't a valid Serializable object"
            logger.error(message)
            raise Exception(message)

        with open(path_dict, "rb") as f:
            d_loaded = load_caller.load(f)

        if len(d_loaded["__serialized_dfs__"]):
            cls._load_dfs(
                folder_path=folder_path,
                dfs=d_loaded["__serialized_dfs__"],
                delete=delete,
                **df_kwargs,
            )

    if delete or delete_only:
        if os.path.isdir(folder_path):
            logger.info("Removing existing directory " + folder_path)
            shutil.rmtree(folder_path)

    if delete_only:
        return None

    obj = cls.from_dict(d_loaded, init_kwargs=init_kwargs)
    return obj

load_any classmethod

load_any(
    path: str = "",
    delete: bool = False,
    delete_only: bool = False,
)

Pass the root path of a serializable object and this will figure out what it is (from the suffix) and call that object's loader function.

For cases where multiple objects work, and I don't want to
have to do if/else logic by suffix

Parameters:

Name	Type	Description	Default
path	str	Path of object. Can exclude the suffix.	''
delete	bool	Delete after load? The default is False.	False
delete_only	bool	Don't load, just delete. The default is False.	False

Returns:

Type	Description
Serializable object of any type

Source code in src\survey_kit\serializable.py

@classmethod
def load_any(cls, path: str = "", delete: bool = False, delete_only: bool = False):
    """
    Pass the root path of a serializable object
        and this will figure out what it is (from the suffix)
        and call that object's loader function.

        For cases where multiple objects work, and I don't want to
        have to do if/else logic by suffix

    Parameters
    ----------
    path : str, optional
        Path of object.  Can exclude the suffix.
    delete : bool, optional
        Delete after load? The default is False.
    delete_only : bool, optional
        Don't load, just delete. The default is False.

    Returns
    -------
    Serializable object of any type

    """

    items = []

    for classi in items:
        suffixi = classi._save_suffix
        if path.endswith(suffixi) or os.path.isdir(f"{path}.{suffixi}"):
            return classi.load(path=path, delete=delete, delete_only=delete_only)

save

save(path: str, quietly: bool = True) -> None

Save a serializable object to

Parameters:

Name	Type	Description	Default
path	str	Path of object. Can exclude the suffix.	required
quietly	bool	No message to console/log	True

Source code in src\survey_kit\serializable.py

def save(self, path: str, quietly: bool = True) -> None:
    """
    Save a serializable object to

    Parameters
    ----------
    path : str, optional
        Path of object.  Can exclude the suffix.
    quietly: bool, optional
        No message to console/log
    """

    self.__serializable_n_dfs = 0

    d_save = self.to_dict()

    if type(path) is not str:
        path = str(path)
    if not path.endswith(f".{self._save_suffix}"):
        folder_path = f"{path}.{self._save_suffix}"
    else:
        folder_path = path

    #   logger.info(folder_path)

    if os.path.isdir(folder_path):
        logger.info("Removing existing directory " + folder_path)
        shutil.rmtree(folder_path)

    #   Make the path to save everything
    create_folders_if_needed([folder_path])
    #   os.makedirs(folder_path)

    #   Save the data
    dfs = d_save["__serialized_dfs__"]
    if len(dfs):
        self._save_dfs(folder_path=folder_path, dfs=dfs, quietly=quietly)

    if self.__fully_serializable__:
        suffix = "json"
        save_caller = json
        w = "w"
    else:
        suffix = "pkl"
        save_caller = pickle
        w = "wb"

    path_object = f"{folder_path}/{self._save_suffix}.{suffix}"
    with open(path_object, w) as f:
        save_caller.dump(d_save, f)

DataFrameList

Bases: Serializable

A light wrapper around Lazy/DataFrames that can be used like a list, mostly, (append, extend, +) and also can be used like the underlying LazyFrame or DataFrame where a narwhals operation is applied to all items in the list.

Examples:

>>> df_list = DataFrameList([df1, df2])
>>> df_list = df_list.filter(nw.col("a") == 1)

This would apply the filter to df1 and df2 and return a DataFrameList with the filtered data.

Source code in src\survey_kit\utilities\dataframe_list.py

class DataFrameList(Serializable):
    """
        A light wrapper around Lazy/DataFrames
        that can be used like a list, mostly, (append, extend, +)
        and also can be used like the underlying LazyFrame or DataFrame
        where a narwhals operation is applied to all items in the list.

        Examples
        --------
        >>> df_list = DataFrameList([df1, df2])
        >>> df_list = df_list.filter(nw.col("a") == 1)

        This would apply the filter to df1 and df2 and return a DataFrameList with
        the filtered data.

    """

    _save_suffix = "df_list"

    def __init__(self, df_list: list[IntoFrameT]):
        self._df_list = df_list

    def __getitem__(self, index):
        return self._df_list[index]

    def __setitem__(self, index, value):
        self._df_list[index] = value

    def __len__(self):
        return len(self._df_list)

    def append(self, df: IntoFrameT):
        self._df_list.append(df)

    def extend(self, iterable: DataFrameList | list[IntoFrameT]):
        if type(iterable) is DataFrameList:
            self._df_list.extend(iterable._df_list)
        else:
            self._df_list.extend(iterable)

    def __add__(self, other: DataFrameList | list[IntoFrameT]):
        if type(other) is DataFrameList:
            self._df_list = self._df_list + other._df_list
            return self
        else:
            self._df_list = self._df_list + other
            return self

    def __repr__(self):
        #   Mimic how lists of polars tables are displayed
        return ",\n".join([dfi.__repr__() for dfi in self._df_list])

    def __iter__(self):
        return iter(self._df_list)

    def __getattr__(self, attr):
        """
        Delegate attribute to polars LazyFrame | DataFrame
            Determines lazy/data based on df_list object 0

        """

        if hasattr(nw.from_native(self._df_list[0]), attr):
            this_attr = getattr(nw.from_native(self._df_list[0]), attr)
            if callable(this_attr):

                def wrapper(*args, **kwargs):
                    output = [
                        getattr(nw.from_native(dfi), attr)(*args, **kwargs)
                        for dfi in self._df_list
                    ]
                    if isinstance(output[0], (nw.LazyFrame, nw.DataFrame)):
                        return DataFrameList([dfi.to_native() for dfi in output])
                    else:
                        return DataFrameList([itemi for itemi in output])

                return wrapper
            else:
                output = [getattr(nw.from_native(dfi), attr) for dfi in self._df_list]
                if isinstance(output[0], (nw.LazyFrame, nw.DataFrame)):
                    return DataFrameList([dfi.to_native() for dfi in output])
                else:
                    return DataFrameList([itemi for itemi in output])

        else:
            raise AttributeError(
                f"{type(nw.from_native(self._df_list[0]))} has no attribute '{attr}'"
            )

    def join_to_list(
        self,
        df_join: list[IntoFrameT],
        on: list[str],
        how: str,
        suffixes: list[str] | None = None,
        prefixes: list[str] | None = None,
    ) -> DataFrameList:
        prefixes = list_input(prefixes)
        suffixes = list_input(suffixes)

        if len(prefixes) and len(prefixes) == len(df_join):
            prefixes = [""] + prefixes
        if len(suffixes) and len(suffixes) == len(df_join):
            suffixes = [""] + suffixes

        for i in range(0, len(self._df_list)):
            self._df_list[i] = join_list(
                [self._df_list[i]] + df_join,
                on=on,
                how=how,
                prefixes=prefixes,
                suffixes=suffixes,
            )
        return self

    def append_list(self, df_append: list[IntoFrameT]) -> DataFrameList:
        for i in range(0, len(self._df_list)):
            self._df_list[i] = concat_wrapper(
                df_list=[self._df_list[i]] + df_append, how="diagonal"
            )

        return self

    def calculate_stats(
        self,
        statistics: list[Statistics] | Statistics | None = None,
        weight: str = "",
        scale_wgts_to: float = 0.0,
        summarize_by: dict[str, list[str]] | None = None,
        display: bool = True,
        display_all_vars: bool = True,
        display_max_vars: int = 20,
        round_output: bool | int = True,
    ) -> DataFrameList:
        def stats_for_one(df: IntoFrameT):
            sc = StatCalculator(
                df=df,
                statistics=statistics,
                weight=weight,
                scale_wgts_to=scale_wgts_to,
                summarize_by=summarize_by,
                display=display,
                display_all_vars=display_all_vars,
                display_max_vars=display_max_vars,
                round_output=round_output,
            )

            return sc.df_estimates

        return self.pipe(stats_for_one)

    def calculate_stats_average(
        self,
        statistics: list[Statistics] | Statistics | None = None,
        weight: str = "",
        scale_wgts_to: float = 0.0,
        summarize_by: dict[str, list[str]] | None = None,
        display: bool = True,
        display_all_vars: bool = True,
        display_max_vars: int = 20,
        round_output: bool | int = True,
    ) -> StatCalculator:
        df_list = self.calculate_stats(
            statistics=statistics,
            weight=weight,
            scale_wgts_to=scale_wgts_to,
            summarize_by=summarize_by,
            display=False,
            display_all_vars=display_all_vars,
            display_max_vars=display_max_vars,
            round_output=round_output,
        )

        sc = StatCalculator(
            df=self[0],
            statistics=statistics,
            weight=weight,
            scale_wgts_to=scale_wgts_to,
            summarize_by=summarize_by,
            display=False,
            display_all_vars=display_all_vars,
            display_max_vars=display_max_vars,
            round_output=round_output,
            calculate=False,
        )
        sc.df_estimates = df_list.average()

        if display:
            sc.print()

        return sc

    def average(self, order_by: list[str] | str | None = None) -> IntoFrameT:
        index = "__df_average_row_index__"
        order_by = list_input(order_by)
        if len(order_by) == 0:
            df_list = [
                nw.from_native(dfi).lazy().collect().with_row_index(name=index)
                for dfi in self._df_list
            ]
        else:
            df_list = [
                nw.from_native(dfi).lazy().with_row_index(name=index, order_by=order_by)
                for dfi in self._df_list
            ]
        df = concat_wrapper(df_list, how="diagonal")

        all_cols = df.collect_schema().names()

        order_by_full = order_by + [index]
        cols_numeric = (
            df.select(cs.numeric() | cs.boolean())
            .drop(order_by + [index])
            .collect_schema()
            .names()
        )
        cols_first = columns_from_list(
            df, columns="*", exclude=cols_numeric + order_by_full
        )

        with_agg = []
        if len(cols_first):
            with_agg.append(nw.col(cols_first).first())
        with_agg.append(nw.col(cols_numeric).mean())
        df_out = nw.from_native(df).lazy().group_by(order_by_full).agg(with_agg)

        return df_out.select(all_cols).sort(order_by_full).drop(index).to_native()

compress_df

compress_df(
    df: IntoFrameT,
    cols: list[str] | str | None = None,
    check_string: bool = False,
    check_string_only: bool = False,
    cast_all_null_to_boolean: bool = True,
    check_date_time: bool = True,
    no_boolean: bool = False,
) -> IntoFrameT

Optimize DataFrame by downcasting numeric types to smallest possible representation.

Analyzes numeric columns and casts them to the smallest data type that can accommodate all values, reducing memory usage and file sizes. Has some optional parameters to handle figuring out the optimal compress for stata

Parameters:

Name	Type	Description	Default
df	IntoFrameT	Input data	required
cols	list[str]	Specific columns to compress (default: all)	None
check_string	bool	Attempt to convert string columns to numeric	False
check_string_only	bool	Only check string conversions	False
cast_all_null_to_boolean	bool	Cast all-null columns to boolean	True
check_date_time	bool	Optimize datetime columns	True
no_boolean	bool	Skip boolean type casting (and leave as int8)	False

Returns:

Type	Description
IntoFrameT	Compressed DataFrame with optimized data types

Examples:

Basic compression:

>>> compressed_df = compress_df(df)

String to numeric conversion:

>>> compressed_df = compress_df(df, check_string=True)

Notes

Automatically detects the smallest integer type that can hold all values in each column, considering ranges like Int8 (-128 to 127), Int16, etc.

Source code in src\survey_kit\utilities\compress.py

def compress_df(
    df: IntoFrameT,
    cols: list[str] | str | None = None,
    check_string: bool = False,
    check_string_only: bool = False,
    cast_all_null_to_boolean: bool = True,
    check_date_time: bool = True,
    no_boolean: bool = False,
) -> IntoFrameT:
    """
    Optimize DataFrame by downcasting numeric types to smallest possible representation.

    Analyzes numeric columns and casts them to the smallest data type that can
    accommodate all values, reducing memory usage and file sizes.
    Has some optional parameters to handle figuring out the optimal compress
    for stata

    Parameters
    ----------
    df : IntoFrameT
        Input data
    cols : list[str], optional
        Specific columns to compress (default: all)
    check_string : bool
        Attempt to convert string columns to numeric
    check_string_only : bool
        Only check string conversions
    cast_all_null_to_boolean : bool
        Cast all-null columns to boolean
    check_date_time : bool
        Optimize datetime columns
    no_boolean : bool
        Skip boolean type casting (and leave as int8)

    Returns
    -------
    IntoFrameT
        Compressed DataFrame with optimized data types

    Examples
    --------
    Basic compression:

    >>> compressed_df = compress_df(df)

    String to numeric conversion:

    >>> compressed_df = compress_df(df, check_string=True)

    Notes
    -----
    Automatically detects the smallest integer type that can hold all values
    in each column, considering ranges like Int8 (-128 to 127), Int16, etc.
    """

    nw_type = NarwhalsType(df)
    df = nw_type.to_polars()

    if check_date_time:
        df = _compress_datetime(df)

    #   Convert numerics
    intlist = {}
    if not no_boolean:
        intlist[pl.Boolean] = [0, 1]
    intlist[pl.Int8] = [-(2**7), 2**7 - 1]
    intlist[pl.Int16] = [-(2**15), 2**15 - 1]
    intlist[pl.Int32] = [-(2**31), 2**31 - 1]
    intlist[pl.Int64] = [-(2**63), 2**63 - 1]

    if cols is None:
        cols = df.lazy().collect_schema().names()

    for columni in cols:
        cast_complete = False
        plType = df.lazy().collect_schema()[columni]

        df_col = df.select(pl.col(columni))

        check_integers = False
        check_float32 = False

        maxValue = None
        minValue = None

        if check_string and (plType == pl.Utf8 or plType == pl.String):
            numeric_string = False
            try:
                df_col = (
                    df_col.select(
                        pl.col(columni).str.strip_chars().cast(pl.Float64, strict=True)
                    )
                    .lazy()
                    .collect()
                )
                numeric_string = True
            except:
                pass

            if numeric_string:
                plType = pl.Float64
                df = df.with_columns(df_col[columni].cast(plType).alias(columni))

        if plType == pl.Float64:
            check_integers = True
            check_float32 = False

            plType_intsize = 65
        elif plType == pl.Float32:
            check_integers = True
            check_float32 = False

            plType_intsize = 65

        elif plType == pl.Int64 or plType == pl.UInt64:
            check_integers = True

            plType_intsize = 64
        elif plType == pl.Int32 or plType == pl.UInt32:
            check_integers = True

            plType_intsize = 32
        elif plType == pl.Int16 or plType == pl.UInt16:
            check_integers = True

            plType_intsize = 16
        elif plType == pl.Int8 or plType == pl.UInt8:
            check_integers = True

            plType_intsize = 8

        #   First check integers
        if check_integers and not check_string_only:
            df_col = df_col.lazy().collect()

            dfCastCheck = df_col.filter(pl.col(columni).is_not_null())

            if dfCastCheck.height == 0 and df_col.height != 0:
                #   All missing, cast to bool (for later combinations to ignore)?
                if cast_all_null_to_boolean:
                    try:
                        #   Try casting on the non-null values
                        dfCastCheck = dfCastCheck.select(
                            pl.col(columni).cast(pl.Boolean, strict=True)
                        )
                        dfCastCheck = None
                        #   Worked - then we're good to do on all of them
                        dfcast = df_col.select(
                            pl.col(columni).cast(pl.Boolean, strict=True)
                        )
                        # logger.info("     Cast " + columni + " as " + str(inti))
                        cast_complete = True
                    except:
                        pass
                        #   logger.warning("     Cannot cast " + columni + " as " + str(pl.Boolean))
            else:
                #   All integers?
                if plType == pl.Float32 or plType == pl.Float64:
                    bAllIntegers = (
                        dfCastCheck.with_columns(pl.col(columni).mod(1) == 0).sum()[
                            0, 0
                        ]
                        == dfCastCheck.height
                    )
                else:
                    bAllIntegers = True

                #   Only downcast to an integer if all are integers
                if bAllIntegers:
                    maxValue = dfCastCheck.max().row(0)[0]
                    minValue = dfCastCheck.min().row(0)[0]

                    if minValue is not None:
                        for inti in intlist:
                            if inti == pl.Boolean:
                                intSize = 1
                            else:
                                intSize = int(str(inti).replace("Int", ""))

                            if plType_intsize > intSize:
                                #   logger.info(intlist[inti])
                                lowerbound = intlist[inti][0]
                                upperbound = intlist[inti][1]
                                #   logger.info(lowerbound)
                                #   logger.info(upperbound)

                                if maxValue <= upperbound and minValue >= lowerbound:
                                    #   logger.info("in range for " + str(inti))

                                    try:
                                        #   Try casting on the non-null values
                                        dfCastCheck = dfCastCheck.select(
                                            pl.col(columni).cast(inti, strict=True)
                                        )
                                        dfCastCheck = None
                                        #   Worked - then we're good to do on all of them
                                        dfcast = df_col.select(
                                            pl.col(columni).cast(inti, strict=True)
                                        )
                                        # logger.info("     Cast " + columni + " as " + str(inti))
                                        cast_complete = True

                                    except:
                                        logger.warning(
                                            "     Cannot cast "
                                            + columni
                                            + " as "
                                            + str(inti)
                                        )

                                    if cast_complete:
                                        break

        if not cast_complete and check_float32 and minValue is not None:
            df_col = df_col.lazy().collect()
            try:
                dfcast = df_col.select(pl.col(columni).cast(pl.Float32, strict=True))
                # logger.info("     Cast " + columni + " as " + str(pl.Float32))
                cast_complete = True
            except:
                logger.warning("     Cannot cast " + columni + " as " + str(pl.Float32))

        if cast_complete:
            df = df.with_columns(dfcast[columni].alias(columni))

    df = nw_type.from_polars(df)

    return NarwhalsType.return_df(df, nw_type)

Config

Global configuration for survey-kit.

Config manages package-wide settings including paths, CPU limits, memory settings, and environment variables. Settings can be configured via environment variables or by directly setting attributes on the config instance.

The config instance is typically accessed via:

    from survey_kit import config
    config.data_root = "/path/to/data"
    config.cpus = 8

Attributes:

Name	Type	Description
code_root	str	Root directory for code files. Set via environment variable _survey_kit_code_root_ or directly. Default is "".
data_root	str	Root directory for data files. Set via environment variable _survey_kit_data_root_ or directly. Default is "".
cpus	int	Number of CPUs to use for parallel operations. Automatically sets thread limits for Polars, OpenBLAS, MKL, etc. when changed. Set via _survey_kit_n_cpus_ or directly. Default is os.cpu_count().
path_temp_files	str	Directory for temporary files. Set via _survey_kit_path_temp_files_ or directly. Default is {data_root}/temp_files.
ram	int	Total available RAM in bytes. Set via _survey_kit_ram_. Default is system total memory.
mem_in_gb	int	Available memory in gigabytes (read-only).
mem_in_mb	int	Available memory in megabytes (read-only).
mem_in_kb	int	Available memory in kilobytes (read-only).

Examples:

Basic configuration:

>>> from survey_kit import config
>>> config.data_root = "/projects/data/myproject"
>>> config.cpus = 16
>>> print(config.path_temp_files)
'/projects/data/myproject/temp_files'

Using environment variables:

    export _survey_kit_data_root_="/projects/data/myproject"
    export _survey_kit_n_cpus_=16

Memory information:

>>> print(f"Available RAM: {config.mem_in_gb} GB")
>>> print(f"Available RAM: {config.mem_in_mb} MB")

Temporary files:

>>> temp_path = config.path_temp_with_random(as_parquet=True)
>>> print(temp_path)
'/projects/data/myproject/temp_files/abc123xyz.parquet'

Notes

Setting cpus automatically updates thread limits for multiple libraries: - POLARS_MAX_THREADS - OMP_NUM_THREADS - NUMEXPR_NUM_THREADS - MKL_NUM_THREADS - OPENBLAS_NUM_THREADS

Source code in src\survey_kit\orchestration\config.py

class Config:
    """
    Global configuration for survey-kit.

    Config manages package-wide settings including paths, CPU limits, memory settings,
    and environment variables. Settings can be configured via environment variables
    or by directly setting attributes on the config instance.

    The config instance is typically accessed via:
    ```python
        from survey_kit import config
        config.data_root = "/path/to/data"
        config.cpus = 8
    ```

    Attributes
    ----------
    code_root : str
        Root directory for code files. Set via environment variable 
        `_survey_kit_code_root_` or directly. Default is "".
    data_root : str
        Root directory for data files. Set via environment variable
        `_survey_kit_data_root_` or directly. Default is "".
    cpus : int
        Number of CPUs to use for parallel operations. Automatically sets
        thread limits for Polars, OpenBLAS, MKL, etc. when changed.
        Set via `_survey_kit_n_cpus_` or directly. Default is os.cpu_count().
    path_temp_files : str
        Directory for temporary files. Set via `_survey_kit_path_temp_files_`
        or directly. Default is {data_root}/temp_files.
    ram : int
        Total available RAM in bytes. Set via `_survey_kit_ram_`.
        Default is system total memory.
    mem_in_gb : int
        Available memory in gigabytes (read-only).
    mem_in_mb : int
        Available memory in megabytes (read-only).
    mem_in_kb : int
        Available memory in kilobytes (read-only).
    Examples
    --------
    Basic configuration:

    >>> from survey_kit import config
    >>> config.data_root = "/projects/data/myproject"
    >>> config.cpus = 16
    >>> print(config.path_temp_files)
    '/projects/data/myproject/temp_files'

    Using environment variables:
    ```bash
        export _survey_kit_data_root_="/projects/data/myproject"
        export _survey_kit_n_cpus_=16
    ```

    Memory information:

    >>> print(f"Available RAM: {config.mem_in_gb} GB")
    >>> print(f"Available RAM: {config.mem_in_mb} MB")

    Temporary files:

    >>> temp_path = config.path_temp_with_random(as_parquet=True)
    >>> print(temp_path)
    '/projects/data/myproject/temp_files/abc123xyz.parquet'

    Notes
    -----
    Setting `cpus` automatically updates thread limits for multiple libraries:
    - POLARS_MAX_THREADS
    - OMP_NUM_THREADS
    - NUMEXPR_NUM_THREADS
    - MKL_NUM_THREADS
    - OPENBLAS_NUM_THREADS

    """

    # parameter_files : dict
    #     Dictionary of parameter file paths. Set via `_survey_kit_parameter_files_`.
    #     Default is {}.
    # pbs_log_path : str
    #     Path for PBS job logs. Set via `_survey_kit_pbs_log_path_`.
    #     Default is "".
    # versions : list
    #     (IGNORE - INTENDED FOR FUTURE USE)
    #     List of version strings (e.g., ["V3", "V2"]). Used to construct
    #     versioned data paths. Set via `_survey_kit_versions_`. Default is [].
    # latest_version : str
    #     (IGNORE - INTENDED FOR FUTURE USE)
    #     Most recent version string from versions list (read-only).
    # data_with_version : str
    #     (IGNORE - INTENDED FOR FUTURE USE)
    #     Path combining data_root with latest_version (read-only).

    _code_root_key = "_survey_kit_code_root_"
    _data_root_key = "_survey_kit_data_root_"
    _version_key = "_survey_kit_versions_"
    _cpus_key = "_survey_kit_n_cpus_"
    _path_temp_files_key = "_survey_kit_path_temp_files_"
    _ram_key = "_survey_kit_ram_"
    _parameter_files_key = "_survey_kit_parameter_files_"
    _pbs_log_path_key = "_survey_kit_pbs_log_path_"

    code_root = TypedEnvVar(_code_root_key, default="", convert=str)
    data_root = TypedEnvVar(_data_root_key, default="", convert=str)
    versions = TypedEnvVar(_version_key, default=[], convert=list)
    _cpus = TypedEnvVar(_cpus_key, os.cpu_count(), int)
    _path_temp_files = TypedEnvVar(_path_temp_files_key, "", str)
    ram = TypedEnvVar(_ram_key, psutil.virtual_memory().total)
    parameter_files = TypedEnvVar(_parameter_files_key, {}, convert=dict)
    pbs_log_path = TypedEnvVar(_pbs_log_path_key, "", str)

    @property
    def latest_version(self) -> str:
        versions = self.versions

        if len(versions):
            return versions[0]

        return ""

    @property
    def data_with_version(self) -> str:
        versions = self.versions

        output = self.data_root
        if len(versions):
            latest = str(versions[0])
            output = os.path.join(output, latest)

        return output

    @property
    def cpus(self) -> int:
        return self._cpus

    @cpus.setter
    def cpus(self, value: int):
        self._cpus = value

        self._set_thread_limits()

    @property
    def path_temp_files(self) -> int:
        if self._path_temp_files != "":
            return self._path_temp_files
        else:
            if self.data_root == "":
                from .. import logger

                message = "You must set Configs().data_root to get a default temp file directory"
                logger.error(message)
                raise Exception(message)

            return Path(self.data_root) / "temp_files"

    @path_temp_files.setter
    def path_temp_files(self, value: str):
        self._path_temp_files = value

    def _set_thread_limits(self):
        n_cpus = self.cpus

        cpu_limits = [
            "POLARS_MAX_THREADS",
            "OMP_NUM_THREADS",
            "NUMEXPR_NUM_THREADS",
            "MKL_NUM_THREADS",
            "OPENBLAS_NUM_THREADS",
        ]
        for limiti in cpu_limits:
            os.environ[limiti] = str(n_cpus)

    @property
    def mem_in_gb(self) -> int:
        return self._mem("gb")

    @property
    def mem_in_mb(self) -> int:
        return self._mem("mb")

    @property
    def mem_in_kb(self) -> int:
        return self._mem("kb")

    def _mem(self, unit: str) -> int:
        unit = unit.lower()

        if unit == "gb":
            power = 3
        elif unit == "mb":
            power = 2
        elif unit == "kb":
            power = 1
        else:
            from .. import logger

            message = f"Must pass kb, mb, or gb ({unit})"
            from .. import logger
            logger.error(message)
            raise Exception(message)

        return int(self.ram / 1024**power)

    def path_temp_with_random(
        self, as_parquet: bool = False, underscore_prefix: bool = False
    ) -> str:
        """
        Generate a random temporary file path.

        Parameters
        ----------
        as_parquet : bool, optional
            Add .parquet extension. Default is False.
        underscore_prefix : bool, optional
            Prefix filename with underscore. Default is False.

        Returns
        -------
        str
            Full path to a uniquely-named temporary file.

        Examples
        --------
        >>> config.path_temp_with_random()
        '/data/temp_files/abc123xyz'

        >>> config.path_temp_with_random(as_parquet=True)
        '/data/temp_files/abc123xyz.parquet'
        """
        if as_parquet:
            parquet_suffix = ".parquet"
        else:
            parquet_suffix = ""

        if underscore_prefix:
            prefix = "_"
        else:
            prefix = ""

        return os.path.normpath(
            f"{self.path_temp_files}/{prefix}{next(tempfile._get_candidate_names())}{parquet_suffix}"
        )


    def clean_temp_directory(self,clean_older_than_days: int = 7):
        from .. import logger

        try:
            path_to_clean = self.path_temp_files

            CleanTempDirectory.clean_old_files(
                temp_dir_path=path_to_clean,
                clean_older_than_days=clean_older_than_days
            )
        except:
            from .. import logger
            logger.info("Not temp directory to clean")

path_temp_with_random

path_temp_with_random(
    as_parquet: bool = False,
    underscore_prefix: bool = False,
) -> str

Generate a random temporary file path.

Parameters:

Name	Type	Description	Default
as_parquet	bool	Add .parquet extension. Default is False.	False
underscore_prefix	bool	Prefix filename with underscore. Default is False.	False

Returns:

Type	Description
str	Full path to a uniquely-named temporary file.

Examples:

>>> config.path_temp_with_random()
'/data/temp_files/abc123xyz'

>>> config.path_temp_with_random(as_parquet=True)
'/data/temp_files/abc123xyz.parquet'

Source code in src\survey_kit\orchestration\config.py

def path_temp_with_random(
    self, as_parquet: bool = False, underscore_prefix: bool = False
) -> str:
    """
    Generate a random temporary file path.

    Parameters
    ----------
    as_parquet : bool, optional
        Add .parquet extension. Default is False.
    underscore_prefix : bool, optional
        Prefix filename with underscore. Default is False.

    Returns
    -------
    str
        Full path to a uniquely-named temporary file.

    Examples
    --------
    >>> config.path_temp_with_random()
    '/data/temp_files/abc123xyz'

    >>> config.path_temp_with_random(as_parquet=True)
    '/data/temp_files/abc123xyz.parquet'
    """
    if as_parquet:
        parquet_suffix = ".parquet"
    else:
        parquet_suffix = ""

    if underscore_prefix:
        prefix = "_"
    else:
        prefix = ""

    return os.path.normpath(
        f"{self.path_temp_files}/{prefix}{next(tempfile._get_candidate_names())}{parquet_suffix}"
    )

random

RandomData

Generate random data, in a slightly easier way.

Parameters:

Name	Type	Description	Default
n_rows	int	Number of rows in the data to be generated	required
seed	int	For replicability, set the seed. Default is 0 (which does not set the seed)	0

Examples:

Create a dataframe with random variables:

>>> nRows = 10000
>>> df = (RandomData(n_rows=nRows, seed=89465551)
...       .index("index")
...       .integer("year", lower=2015, upper=2020)
...       .float("var1", 0, 100000)
...       .integer("var2", 1, 100000)
...       .integer("var3", 1, 50)
...       .float("var4", 0, 1)
...       .date("date", date(2020, 1, 1), date(2025, 12, 31))
...       .datetime("datetime", date(2020, 1, 1), date(2025, 12, 31))
...       .np_distribution("v_normal", "normal", dict(loc=1, scale=2))
...       .np_distribution("v_lognormal", "lognormal", dict(mean=1, sigma=2))
...       .to_df()
... )

Source code in src\survey_kit\utilities\random.py

class RandomData:
    """
    Generate random data, in a slightly easier way.

    Parameters
    ----------
    n_rows : int
        Number of rows in the data to be generated
    seed : int, optional
        For replicability, set the seed.
        Default is 0 (which does not set the seed)

    Examples
    --------
    Create a dataframe with random variables:

    >>> nRows = 10000
    >>> df = (RandomData(n_rows=nRows, seed=89465551)
    ...       .index("index")
    ...       .integer("year", lower=2015, upper=2020)
    ...       .float("var1", 0, 100000)
    ...       .integer("var2", 1, 100000)
    ...       .integer("var3", 1, 50)
    ...       .float("var4", 0, 1)
    ...       .date("date", date(2020, 1, 1), date(2025, 12, 31))
    ...       .datetime("datetime", date(2020, 1, 1), date(2025, 12, 31))
    ...       .np_distribution("v_normal", "normal", dict(loc=1, scale=2))
    ...       .np_distribution("v_lognormal", "lognormal", dict(mean=1, sigma=2))
    ...       .to_df()
    ... )
    """

    def __init__(self, n_rows: int, seed: int = 0):
        if seed > 0:
            set_seed(seed)

        self.rng = RandomNumberGenerator()
        self.n_rows = n_rows
        self._data = {}

    def index(self, name: str) -> RandomData:
        """
        Create an index column (i.e. 0-n_rows-1)

        Parameters
        ----------
        name : str
            column name

        Returns
        -------
        RandomData object (so you can chain 's)
        """
        self._data[name] = range(0, self.n_rows)

        return self

    def boolean(self, name: str) -> RandomData:
        """
        Create an boolean column

        Parameters
        ----------
        name : str
            column name

        Returns
        -------
        RandomData object (so you can chain 's)
        """

        self._data[name] = np.ceil(self.rng.uniform(low=-1, high=1, size=self.n_rows))

        return self

    def integer(self, name: str, lower: int, upper: int) -> RandomData:
        """
        Create an integer column in [lower,upper]

        Parameters
        ----------
        name : str
            column name
        lower : int
            lower bound (inclusive)
        upper : int
            upper bound (inclusive)

        Returns
        -------
        RandomData object (so you can chain 's)
        """

        self._data[name] = np.ceil(
            self.rng.uniform(low=lower - 1, high=upper, size=self.n_rows)
        )

        return self

    def float(self, name: str, lower: float, upper: float) -> RandomData:
        """
        Create a float64 column in (lower,upper)

        Parameters
        ----------
        name : str
            column name
        lower : float
            lower bound
        upper : float
            upper bound

        Returns
        -------
        RandomData object (so you can chain 's)
        """

        self._data[name] = self.rng.uniform(low=lower, high=upper, size=self.n_rows)

        return self

    def date(self, name: str, start: date, end: date) -> RandomData:
        """
        Create a date column in [start,end]

        Parameters
        ----------
        name : str
            column name
        start : date
            lower bound
        end : date
            upper bound

        Returns
        -------
        RandomData object (so you can chain 's)
        """
        self._data[name] = pl.date_range(start, end, "1d", eager=True).sample(
            n=self.n_rows, with_replacement=True
        )

        return self

    def datetime(
        self, name: str, start: datetime | date, end: datetime | date
    ) -> RandomData:
        """
        Create a datetime column in [start,end]

        Parameters
        ----------
        name : str
            column name
        start : datetime | date
            lower bound
        end : datetime | date
            upper bound

        Returns
        -------
        RandomData object (so you can chain 's)
        """

        self._data[name] = pl.datetime_range(start, end, "1m", eager=True).sample(
            n=self.n_rows, with_replacement=True
        )

        return self

    def np_distribution(self, name: str, distribution: str, **kwargs) -> RandomData:
        """
        Create a column from a numpy distribution.

        Parameters
        ----------
        name : str
            Column name.
        distribution : str
            Name of numpy random distribution (e.g., "normal", "lognormal", "exponential").
            Must be a valid method of np.random.Generator.
        **kwargs
            Keyword arguments passed to the distribution function (e.g., loc, scale for normal).

        Returns
        -------
        RandomData
            Self for method chaining.

        Raises
        ------
        Exception
            If distribution is not a valid numpy random distribution.

        Examples
        --------
        >>> df = (RandomData(n_rows=1000, seed=123)
        ...       .np_distribution("normal_var", "normal", loc=0, scale=1)
        ...       .np_distribution("lognormal_var", "lognormal", mean=1, sigma=2)
        ...       .np_distribution("exponential_var", "exponential", scale=1.5)
        ...       .to_df())
        """

        if hasattr(self.rng, distribution):
            generator = getattr(self.rng, distribution)
            self._data[name] = generator(size=self.n_rows, **kwargs)
        else:
            message = f"Numpy generator does not have the function '{distribution}'"
            logger.error(message)
            raise Exception(message)

        return self

    def to_df(self, compress: bool = True) -> pl.DataFrame:
        """
        Convert accumulated random data to a Polars DataFrame.

        Parameters
        ----------
        compress : bool, optional
            Apply compression to reduce memory usage by downcasting numeric types.
            Default is True.

        Returns
        -------
        pl.DataFrame
            DataFrame containing all generated columns.

        Examples
        --------
        >>> df = (RandomData(n_rows=1000, seed=123)
        ...       .integer("id", 1, 1000)
        ...       .float("value", 0, 100)
        ...       .to_df())
        >>> print(df)

        >>> # Without compression
        >>> df_uncompressed = (RandomData(n_rows=1000, seed=123)
        ...                    .integer("id", 1, 1000)
        ...                    .to_df(compress=False))
        """

        df = pl.DataFrame(self._data)
        if compress:
            return compress_df(df)
        else:
            return df

boolean

boolean(name: str) -> RandomData

Create an boolean column

Parameters:

Name	Type	Description	Default
name	str	column name	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def boolean(self, name: str) -> RandomData:
    """
    Create an boolean column

    Parameters
    ----------
    name : str
        column name

    Returns
    -------
    RandomData object (so you can chain 's)
    """

    self._data[name] = np.ceil(self.rng.uniform(low=-1, high=1, size=self.n_rows))

    return self

date

date(name: str, start: date, end: date) -> RandomData

Create a date column in [start,end]

Parameters:

Name	Type	Description	Default
name	str	column name	required
start	date	lower bound	required
end	date	upper bound	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def date(self, name: str, start: date, end: date) -> RandomData:
    """
    Create a date column in [start,end]

    Parameters
    ----------
    name : str
        column name
    start : date
        lower bound
    end : date
        upper bound

    Returns
    -------
    RandomData object (so you can chain 's)
    """
    self._data[name] = pl.date_range(start, end, "1d", eager=True).sample(
        n=self.n_rows, with_replacement=True
    )

    return self

datetime

datetime(
    name: str, start: datetime | date, end: datetime | date
) -> RandomData

Create a datetime column in [start,end]

Parameters:

Name	Type	Description	Default
name	str	column name	required
start	datetime | date	lower bound	required
end	datetime | date	upper bound	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def datetime(
    self, name: str, start: datetime | date, end: datetime | date
) -> RandomData:
    """
    Create a datetime column in [start,end]

    Parameters
    ----------
    name : str
        column name
    start : datetime | date
        lower bound
    end : datetime | date
        upper bound

    Returns
    -------
    RandomData object (so you can chain 's)
    """

    self._data[name] = pl.datetime_range(start, end, "1m", eager=True).sample(
        n=self.n_rows, with_replacement=True
    )

    return self

float

float(name: str, lower: float, upper: float) -> RandomData

Create a float64 column in (lower,upper)

Parameters:

Name	Type	Description	Default
name	str	column name	required
lower	float	lower bound	required
upper	float	upper bound	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def float(self, name: str, lower: float, upper: float) -> RandomData:
    """
    Create a float64 column in (lower,upper)

    Parameters
    ----------
    name : str
        column name
    lower : float
        lower bound
    upper : float
        upper bound

    Returns
    -------
    RandomData object (so you can chain 's)
    """

    self._data[name] = self.rng.uniform(low=lower, high=upper, size=self.n_rows)

    return self

index

index(name: str) -> RandomData

Create an index column (i.e. 0-n_rows-1)

Parameters:

Name	Type	Description	Default
name	str	column name	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def index(self, name: str) -> RandomData:
    """
    Create an index column (i.e. 0-n_rows-1)

    Parameters
    ----------
    name : str
        column name

    Returns
    -------
    RandomData object (so you can chain 's)
    """
    self._data[name] = range(0, self.n_rows)

    return self

integer

integer(name: str, lower: int, upper: int) -> RandomData

Create an integer column in [lower,upper]

Parameters:

Name	Type	Description	Default
name	str	column name	required
lower	int	lower bound (inclusive)	required
upper	int	upper bound (inclusive)	required

Returns:

Type	Description
RandomData object (so you can chain 's)

Source code in src\survey_kit\utilities\random.py

def integer(self, name: str, lower: int, upper: int) -> RandomData:
    """
    Create an integer column in [lower,upper]

    Parameters
    ----------
    name : str
        column name
    lower : int
        lower bound (inclusive)
    upper : int
        upper bound (inclusive)

    Returns
    -------
    RandomData object (so you can chain 's)
    """

    self._data[name] = np.ceil(
        self.rng.uniform(low=lower - 1, high=upper, size=self.n_rows)
    )

    return self

np_distribution

np_distribution(
    name: str, distribution: str, **kwargs
) -> RandomData

Create a column from a numpy distribution.

Parameters:

Name	Type	Description	Default
name	str	Column name.	required
distribution	str	Name of numpy random distribution (e.g., "normal", "lognormal", "exponential"). Must be a valid method of np.random.Generator.	required
**kwargs		Keyword arguments passed to the distribution function (e.g., loc, scale for normal).	{}

Returns:

Type	Description
RandomData	Self for method chaining.

Raises:

Type	Description
Exception	If distribution is not a valid numpy random distribution.

Examples:

>>> df = (RandomData(n_rows=1000, seed=123)
...       .np_distribution("normal_var", "normal", loc=0, scale=1)
...       .np_distribution("lognormal_var", "lognormal", mean=1, sigma=2)
...       .np_distribution("exponential_var", "exponential", scale=1.5)
...       .to_df())

Source code in src\survey_kit\utilities\random.py

def np_distribution(self, name: str, distribution: str, **kwargs) -> RandomData:
    """
    Create a column from a numpy distribution.

    Parameters
    ----------
    name : str
        Column name.
    distribution : str
        Name of numpy random distribution (e.g., "normal", "lognormal", "exponential").
        Must be a valid method of np.random.Generator.
    **kwargs
        Keyword arguments passed to the distribution function (e.g., loc, scale for normal).

    Returns
    -------
    RandomData
        Self for method chaining.

    Raises
    ------
    Exception
        If distribution is not a valid numpy random distribution.

    Examples
    --------
    >>> df = (RandomData(n_rows=1000, seed=123)
    ...       .np_distribution("normal_var", "normal", loc=0, scale=1)
    ...       .np_distribution("lognormal_var", "lognormal", mean=1, sigma=2)
    ...       .np_distribution("exponential_var", "exponential", scale=1.5)
    ...       .to_df())
    """

    if hasattr(self.rng, distribution):
        generator = getattr(self.rng, distribution)
        self._data[name] = generator(size=self.n_rows, **kwargs)
    else:
        message = f"Numpy generator does not have the function '{distribution}'"
        logger.error(message)
        raise Exception(message)

    return self

to_df

to_df(compress: bool = True) -> pl.DataFrame

Convert accumulated random data to a Polars DataFrame.

Parameters:

Name	Type	Description	Default
compress	bool	Apply compression to reduce memory usage by downcasting numeric types. Default is True.	True

Returns:

Type	Description
DataFrame	DataFrame containing all generated columns.

Examples:

>>> df = (RandomData(n_rows=1000, seed=123)
...       .integer("id", 1, 1000)
...       .float("value", 0, 100)
...       .to_df())
>>> print(df)

>>> # Without compression
>>> df_uncompressed = (RandomData(n_rows=1000, seed=123)
...                    .integer("id", 1, 1000)
...                    .to_df(compress=False))

Source code in src\survey_kit\utilities\random.py

def to_df(self, compress: bool = True) -> pl.DataFrame:
    """
    Convert accumulated random data to a Polars DataFrame.

    Parameters
    ----------
    compress : bool, optional
        Apply compression to reduce memory usage by downcasting numeric types.
        Default is True.

    Returns
    -------
    pl.DataFrame
        DataFrame containing all generated columns.

    Examples
    --------
    >>> df = (RandomData(n_rows=1000, seed=123)
    ...       .integer("id", 1, 1000)
    ...       .float("value", 0, 100)
    ...       .to_df())
    >>> print(df)

    >>> # Without compression
    >>> df_uncompressed = (RandomData(n_rows=1000, seed=123)
    ...                    .integer("id", 1, 1000)
    ...                    .to_df(compress=False))
    """

    df = pl.DataFrame(self._data)
    if compress:
        return compress_df(df)
    else:
        return df

RandomNumberGenerator

RandomNumberGenerator() -> np.random.Generator

Create a new numpy random number generator with random seed.

Returns:

Type	Description
Generator	Numpy random number generator initialized with a random seed from Python's random module.

Examples:

>>> from survey_kit.utilities.random import RandomNumberGenerator
>>> rng = RandomNumberGenerator()
>>> random_values = rng.normal(loc=0, scale=1, size=100)

Source code in src\survey_kit\utilities\random.py

def RandomNumberGenerator() -> np.random.Generator:
    """
    Create a new numpy random number generator with random seed.

    Returns
    -------
    np.random.Generator
        Numpy random number generator initialized with a random seed
        from Python's random module.

    Examples
    --------
    >>> from survey_kit.utilities.random import RandomNumberGenerator
    >>> rng = RandomNumberGenerator()
    >>> random_values = rng.normal(loc=0, scale=1, size=100)
    """
    return np.random.default_rng(random.randint(1, 2**63 - 1))

generate_seed

generate_seed(power_of_2_limit: int = 32)

Generate a random seed value.

Parameters:

Name	Type	Description	Default
power_of_2_limit	int	Upper bound is 2^power_of_2_limit. Default is 32.	32

Returns:

Type	Description
int	Random integer between 1 and 2^power_of_2_limit - 1.

Examples:

>>> from survey_kit.utilities.random import generate_seed, set_seed
>>> seed = generate_seed()
>>> set_seed(seed)  # Use for reproducibility

Source code in src\survey_kit\utilities\random.py

def generate_seed(power_of_2_limit: int = 32):
    """
    Generate a random seed value.

    Parameters
    ----------
    power_of_2_limit : int, optional
        Upper bound is 2^power_of_2_limit. Default is 32.

    Returns
    -------
    int
        Random integer between 1 and 2^power_of_2_limit - 1.

    Examples
    --------
    >>> from survey_kit.utilities.random import generate_seed, set_seed
    >>> seed = generate_seed()
    >>> set_seed(seed)  # Use for reproducibility
    """
    rng = RandomNumberGenerator()
    return int(rng.integers(1, 2**power_of_2_limit - 1, 1)[0])

get_random_state

get_random_state()

Get the current state of the random number generator.

Returns:

Type	Description
tuple	Internal state of Python's random module.

Examples:

>>> from survey_kit.utilities.random import get_random_state, set_random_state
>>> state = get_random_state()
>>> # ... do some random operations ...
>>> set_random_state(state)  # Restore to previous state

Source code in src\survey_kit\utilities\random.py

def get_random_state():
    """
    Get the current state of the random number generator.

    Returns
    -------
    tuple
        Internal state of Python's random module.

    Examples
    --------
    >>> from survey_kit.utilities.random import get_random_state, set_random_state
    >>> state = get_random_state()
    >>> # ... do some random operations ...
    >>> set_random_state(state)  # Restore to previous state
    """
    return random.getstate()

set_random_state

set_random_state(value)

Set the random number generator to a specific state.

Parameters:

Name	Type	Description	Default
value	tuple	State tuple from get_random_state().	required

Examples:

>>> state = get_random_state()
>>> # ... do some random operations ...
>>> set_random_state(state)  # Restore to previous state

Source code in src\survey_kit\utilities\random.py

def set_random_state(value):
    """
    Set the random number generator to a specific state.

    Parameters
    ----------
    value : tuple
        State tuple from get_random_state().

    Examples
    --------
    >>> state = get_random_state()
    >>> # ... do some random operations ...
    >>> set_random_state(state)  # Restore to previous state
    """
    random.setstate(value)

set_seed

set_seed(seed: int = 0)

Set the random seed for reproducibility.

Sets the seed for Python's random module, which is used by the random number generator in this module.

Parameters:

Name	Type	Description	Default
seed	int	Random seed. If 0 or less, does not set seed. Default is 0.	0

Examples:

>>> from survey_kit.utilities.random import set_seed, generate_seed
>>> set_seed(12345)
>>> seed1 = generate_seed()
>>> set_seed(12345)
>>> seed2 = generate_seed()
>>> assert seed1 == seed2  # Same seed produces same sequence

Source code in src\survey_kit\utilities\random.py

def set_seed(seed: int = 0):
    """
    Set the random seed for reproducibility.

    Sets the seed for Python's random module, which is used by the
    random number generator in this module.

    Parameters
    ----------
    seed : int, optional
        Random seed. If 0 or less, does not set seed. Default is 0.

    Examples
    --------
    >>> from survey_kit.utilities.random import set_seed, generate_seed
    >>> set_seed(12345)
    >>> seed1 = generate_seed()
    >>> set_seed(12345)
    >>> seed2 = generate_seed()
    >>> assert seed1 == seed2  # Same seed produces same sequence
    """

    if seed > 0:
        random.seed(seed)

FormulaBuilder

Build and manipulate R-style formulas for statistical models.

FormulaBuilder provides a programmatic way to construct complex model formulas using R/Patsy-style syntax. It supports formula manipulation, variable expansion, interactions, transformations, and pattern matching against dataframes.

The class works with Formulaic to parse and expand formulas, and integrates with the dataframe utilities to resolve wildcards and column patterns.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Reference dataframe for resolving column names and wildcards. Default is None.	None
formula	str	Initial formula string. If empty, constructs from lhs and constant. Default is "".	''
lhs	str	Left-hand side of formula (response variable). Default is "".	''
constant	bool	Include intercept term (1) or suppress it (0). Default is True.	True

Attributes:

Name	Type	Description
formula	str	Current formula string.
df	IntoFrameT | None	Reference dataframe.
columns	list[str]	All variables required by the formula.
columns_rhs	list[str]	Variables in right-hand side of formula.
columns_lhs	list[str]	Variables in left-hand side of formula.

Examples:

Basic formula construction:

>>> from survey_kit.utilities.formula import FormulaBuilder
>>> 
>>> fb = FormulaBuilder(df=df, lhs="income", constant=True)
>>> fb += "age + education"
>>> print(fb.formula)
'income~1+age+education'

Add variables using wildcards:

>>> fb = FormulaBuilder(df=df, lhs="income")
>>> fb.continuous(columns="demographic_*")
>>> fb.factor(columns="region")
>>> print(fb.formula)

Polynomial terms:

>>> fb = FormulaBuilder(df=df, lhs="income")
>>> fb.polynomial(columns="age", degree=3)
>>> print(fb.formula)
'income~1+poly(age,degree=3,raw=True)'

Interactions:

>>> fb = FormulaBuilder(df=df, lhs="income")
>>> fb.simple_interaction(columns=["age", "education"], order=2)
>>> print(fb.formula)
'income~1+(age+education)**2'

Standardization:

>>> fb = FormulaBuilder(df=df, lhs="income")
>>> fb.scale(columns=["age", "experience"])
>>> print(fb.formula)
'income~1+scale(age)+scale(experience)'

Working with existing formula strings:

>>> formula = "income~age+education+age:education"
>>> fb = FormulaBuilder(df=df, formula=formula)
>>> fb.expand()  # Expand shorthand
>>> print(fb.rhs())  # Get right-hand side
'age+education+age:education'
>>> print(fb.columns)
['income', 'age', 'education']

Notes

Formula syntax follows R/Patsy conventions: - ~ separates left and right sides - + adds terms - : creates interactions - * creates main effects and interactions: a*b = a+b+a:b - **n creates all n-way interactions - I() for arithmetic operations - C() for categorical variables - Functions like scale(), center(), poly() for transformations

Wildcards in column specifications: - "income_*" matches all columns starting with "income_" - ["age", "education_*"] matches age and all education columns

The FormulaBuilder can be used in two modes: 1. Object mode: Create instance and chain methods 2. Static mode: Call methods with self=None for one-off operations

See Also

formulaic.Formula : Underlying formula parser

Source code in src\survey_kit\utilities\formula_builder.py

class FormulaBuilder:
    """
    Build and manipulate R-style formulas for statistical models.

    FormulaBuilder provides a programmatic way to construct complex model formulas
    using R/Patsy-style syntax. It supports formula manipulation, variable expansion,
    interactions, transformations, and pattern matching against dataframes.

    The class works with Formulaic to parse and expand formulas, and integrates
    with the dataframe utilities to resolve wildcards and column patterns.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Reference dataframe for resolving column names and wildcards.
        Default is None.
    formula : str, optional
        Initial formula string. If empty, constructs from lhs and constant.
        Default is "".
    lhs : str, optional
        Left-hand side of formula (response variable). Default is "".
    constant : bool, optional
        Include intercept term (1) or suppress it (0). Default is True.

    Attributes
    ----------
    formula : str
        Current formula string.
    df : IntoFrameT | None
        Reference dataframe.
    columns : list[str]
        All variables required by the formula.
    columns_rhs : list[str]
        Variables in right-hand side of formula.
    columns_lhs : list[str]
        Variables in left-hand side of formula.

    Examples
    --------
    Basic formula construction:

    >>> from survey_kit.utilities.formula import FormulaBuilder
    >>> 
    >>> fb = FormulaBuilder(df=df, lhs="income", constant=True)
    >>> fb += "age + education"
    >>> print(fb.formula)
    'income~1+age+education'

    Add variables using wildcards:

    >>> fb = FormulaBuilder(df=df, lhs="income")
    >>> fb.continuous(columns="demographic_*")
    >>> fb.factor(columns="region")
    >>> print(fb.formula)

    Polynomial terms:

    >>> fb = FormulaBuilder(df=df, lhs="income")
    >>> fb.polynomial(columns="age", degree=3)
    >>> print(fb.formula)
    'income~1+poly(age,degree=3,raw=True)'

    Interactions:

    >>> fb = FormulaBuilder(df=df, lhs="income")
    >>> fb.simple_interaction(columns=["age", "education"], order=2)
    >>> print(fb.formula)
    'income~1+(age+education)**2'

    Standardization:

    >>> fb = FormulaBuilder(df=df, lhs="income")
    >>> fb.scale(columns=["age", "experience"])
    >>> print(fb.formula)
    'income~1+scale(age)+scale(experience)'

    Working with existing formula strings:

    >>> formula = "income~age+education+age:education"
    >>> fb = FormulaBuilder(df=df, formula=formula)
    >>> fb.expand()  # Expand shorthand
    >>> print(fb.rhs())  # Get right-hand side
    'age+education+age:education'
    >>> print(fb.columns)
    ['income', 'age', 'education']

    Notes
    -----
    Formula syntax follows R/Patsy conventions:
    - `~` separates left and right sides
    - `+` adds terms
    - `:` creates interactions
    - `*` creates main effects and interactions: `a*b` = `a+b+a:b`
    - `**n` creates all n-way interactions
    - `I()` for arithmetic operations
    - `C()` for categorical variables
    - Functions like `scale()`, `center()`, `poly()` for transformations

    Wildcards in column specifications:
    - `"income_*"` matches all columns starting with "income_"
    - `["age", "education_*"]` matches age and all education columns

    The FormulaBuilder can be used in two modes:
    1. Object mode: Create instance and chain methods
    2. Static mode: Call methods with self=None for one-off operations

    See Also
    --------
    formulaic.Formula : Underlying formula parser
    """

    def __init__(
        self,
        df: IntoFrameT | None = None,
        formula: str = "",
        lhs: str = "",
        constant: bool = True,
    ):
        if formula == "":
            self.formula = f"{lhs}~{int(constant)}"
        else:
            self.formula = formula

        if df is not None:
            self.df = nw.from_native(df).head(0).to_native()
        else:
            self.df = None

    def __str__(self):
        return self.formula

    def __add__(self, o):
        """Add terms to the formula using + operator."""
        if type(o) is FormulaBuilder:
            o = o.rhs()

        o = str(o)
        if o.startswith("~"):
            o = o[1 : len(o)]

        if o.startswith("1+"):
            o = o[2 : len(o)]
        if o.startswith("0+"):
            o = o[2 : len(o)]

        self.formula = f"{self.formula}+{o}"

        self.expand()

        return self

    def add_to_formula(self, add_part: str = "", plus_first: bool = True) -> None:
        """
        Append a string to the formula.

        Parameters
        ----------
        add_part : str, optional
            String to append. Default is "".
        plus_first : bool, optional
            Add "+" before the string. Default is True.
        """
        if plus_first:
            plus = "+"
        else:
            plus = ""
        self.formula += f"{plus}{add_part}"

    def any_wrapper(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        case_insensitive: bool = False,
        prefix: str = "",
        suffix: str = "",
    ) -> str | FormulaBuilder:
        """
        General wrapper for adding columns with prefix/suffix.

        Used internally by other methods to add variables with transformations.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for resolving column patterns. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause (bypasses column lookup). Default is "".
        case_insensitive : bool, optional
            Case-insensitive column matching. Default is False.
        prefix : str, optional
            String to prepend to each column. Default is "".
        suffix : str, optional
            String to append to each column. Default is "".

        Returns
        -------
        str | FormulaBuilder
            Formula string if self is None, otherwise returns self for chaining.
        """
        columns = list_input(columns)

        #   Dataframe to look for columns in
        if df is None:
            df = self.df

        if clause != "":
            output = f"+{prefix}{clause}{suffix}"
        else:
            if df is not None:
                columns = columns_from_list(
                    df=df, columns=columns, case_insensitive=case_insensitive
                )

            out_list = [f"{prefix}{coli}{suffix}" for coli in columns]
            output = "+" + "+".join(out_list)

        if self is None:
            return output
        else:
            self.formula += output
            return self

    def continuous(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        case_insensitive: bool = False,
    ) -> str | FormulaBuilder:
        """
        Add continuous variables to the formula.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns (e.g., "income_*"). Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb = FormulaBuilder(df=df, lhs="y")
        >>> fb.continuous(columns=["age", "income"])
        >>> print(fb.formula)
        'y~1+age+income'
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        return caller.any_wrapper(
            df=df, columns=columns, clause=clause, case_insensitive=case_insensitive
        )

    def function(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        operator_before: str = "",
        operator_after: str = "",
        function_item: str = "",
        case_insensitive: bool = False,
        **kwargs,
    ) -> str | FormulaBuilder:
        """
        Wrap columns in a function call.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        operator_before : str, optional
            String before function call. Default is "".
        operator_after : str, optional
            String after function call. Default is "".
        function_item : str, optional
            Function name (e.g., "log", "sqrt"). Default is "".
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.
        **kwargs
            Additional function arguments.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.function(columns="income", function_item="log")
        >>> print(fb.formula)
        'y~1+log(income)'
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        if function_item != "":
            operator_before = f"{operator_before}{function_item}("

            operator_after_final = ""
            if len(kwargs):
                for keyi, valuei in kwargs.items():
                    operator_after_final += ","
                    if type(valuei) is str:
                        operator_after_final += f"{keyi}='{valuei}'"
                    else:
                        operator_after_final += f"{keyi}={valuei}"

            operator_after_final += f"{operator_after})"
        else:
            operator_after_final = operator_after

        return caller.any_wrapper(
            df=df,
            columns=columns,
            clause=clause,
            case_insensitive=case_insensitive,
            prefix=f"{{{operator_before}",
            suffix=f"{operator_after_final}}}",
        )

    def scale(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        standardize: bool = True,
        case_insensitive: bool = False,
    ) -> str | FormulaBuilder:
        """
        Standardize or center variables.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        standardize : bool, optional
            If True, standardize (mean=0, sd=1). If False, only center (mean=0).
            Default is True.
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.scale(columns="age")  # Standardize
        >>> fb.scale(columns="income", standardize=False)  # Center only
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        if standardize:
            function_item = "scale"
        else:
            function_item = "center"

        return caller.any_wrapper(
            df=df,
            columns=columns,
            clause=clause,
            case_insensitive=case_insensitive,
            prefix=f"{function_item}(",
            suffix=")",
        )

    def center(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        case_insensitive: bool = False,
    ) -> str | FormulaBuilder:
        """
        Center variables (subtract mean).

        Convenience method for scale(standardize=False).

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        return caller.scale(
            df=df,
            columns=columns,
            clause=clause,
            case_insensitive=case_insensitive,
            standardize=False,
        )

    def standardize(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        case_insensitive: bool = False,
    ) -> str | FormulaBuilder:
        """
        Standardize variables (mean=0, sd=1).

        Convenience method for scale(standardize=True).

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        return caller.scale(
            df=df,
            columns=columns,
            clause=clause,
            case_insensitive=case_insensitive,
            standardize=True,
        )

    def polynomial(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        degree: int = 0,
        case_insensitive: bool = False,
        center: bool = False,
    ) -> str | FormulaBuilder:
        """
        Add polynomial terms.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        degree : int, optional
            Polynomial degree. Default is 0 (returns continuous).
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.
        center : bool, optional
            Use orthogonal polynomials (centered). Default is False (raw polynomials).

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.polynomial(columns="age", degree=3)
        >>> print(fb.formula)
        'y~1+poly(age,degree=3,raw=True)'

        >>> fb.polynomial(columns="age", degree=2, center=True)
        >>> print(fb.formula)
        'y~1+poly(age,degree=2,raw=False)'
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self
            if df is None:
                df = self.df

        if degree <= 1:
            if center:
                return caller.center(
                    df=df,
                    columns=columns,
                    clause=clause,
                    case_insensitive=case_insensitive,
                )
            else:
                return caller.continuous(
                    df=df,
                    columns=columns,
                    clause=clause,
                    case_insensitive=case_insensitive,
                )
        else:
            subformula = ""

            for power in range(1, degree + 1):
                operator_before = "poly("
                if center:
                    operator_after = f",degree={degree},raw=False)"
                else:
                    operator_after = f",degree={degree},raw=True)"

                subformula += FormulaBuilder.function(
                    df=df,
                    columns=columns,
                    clause=clause,
                    operator_before=operator_before,
                    operator_after=operator_after,
                    case_insensitive=case_insensitive,
                )

            if self is None:
                return subformula
            else:
                self.add_to_formula(subformula, False)
                return self

    def simple_interaction(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        order: int = 2,
        case_insensitive: bool = False,
        sub_function: Callable | None = None,
        no_base: bool = False,
    ) -> str | FormulaBuilder:
        """
        Create interactions between variables.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        order : int, optional
            Interaction order. 2 = pairwise, 3 = three-way, etc. Default is 2.
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.
        sub_function : Callable | None, optional
            Function to apply to columns before interacting. Default is None.
        no_base : bool, optional
            Exclude main effects (only interactions). Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.simple_interaction(columns=["age", "education"], order=2)
        >>> print(fb.formula)
        'y~1+(age+education)**2'

        >>> fb.simple_interaction(columns=["a", "b", "c"], order=2, no_base=True)
        >>> print(fb.formula)
        'y~1+(a+b+c)**2-(a+b+c)'  # Interactions only
        """
        if self is not None:
            if df is None:
                df = self.df

        if sub_function is None:
            sub_function = FormulaBuilder.continuous

        columns = columns_from_list(
            df=df, columns=columns, case_insensitive=case_insensitive
        )

        subformula = sub_function(
            df=df, columns=columns, case_insensitive=case_insensitive
        )
        #   Remove leading plus sign
        subformula = subformula[1 : len(subformula)]

        if len(columns) > 1:
            output = f"({subformula})**{order}"
        else:
            output = subformula

        if no_base:
            output += f"-({subformula})"

        if self is None:
            return f"+{output}"
        else:
            self.add_to_formula(output)
            return self

    def interact_clauses(
        self=None, clause1: str = "", clause2: str = "", no_base: bool = False
    ) -> str | FormulaBuilder:
        """
        Create interactions between two formula clauses.

        Parameters
        ----------
        clause1 : str, optional
            First formula clause. Default is "".
        clause2 : str, optional
            Second formula clause. Default is "".
        no_base : bool, optional
            Exclude main effects from clauses. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.interact_clauses("age+education", "region")
        >>> print(fb.formula)
        'y~1+(age+education)*(region)'
        """
        if clause1.startswith("+"):
            clause1 = clause1[1:]
        if clause2.startswith("+"):
            clause2 = clause2[1:]

        clause1 = clause1.replace("++", "+")
        clause2 = clause2.replace("++", "+")

        output = f"({clause1})*({clause2})"

        if no_base:
            output += f"-({clause1} + {clause2})"

        if self is None:
            return f"+{output}"
        else:
            self.add_to_formula(output)
            return self

    def factor(
        self=None,
        df: IntoFrameT | None = None,
        columns: str | list | None = None,
        clause: str = "",
        reference=None,
        case_insensitive: bool = False,
    ):
        """
        Add categorical variables with treatment coding.

        Parameters
        ----------
        df : IntoFrameT | None, optional
            Dataframe for column lookup. Default is None.
        columns : str | list | None, optional
            Column names or patterns. Default is None.
        clause : str, optional
            Pre-constructed clause. Default is "".
        reference : str | int | None, optional
            Reference level for treatment coding. Default is None (use first level).
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Formula string or self for chaining.

        Examples
        --------
        >>> fb.factor(columns="region")
        >>> print(fb.formula)
        'y~1+C(region)'

        >>> fb.factor(columns="education", reference="high_school")
        >>> print(fb.formula)
        "y~1+C(education, contr.treatment('high_school'))"
        """
        if self is None:
            caller = FormulaBuilder
        else:
            caller = self

        prefix = "C("
        if reference is not None:
            if type(reference) is str:
                suffix = f", contr.treatment('{reference}'))"
            else:
                suffix = f", contr.treatment({reference}))"

        else:
            suffix = ")"  #  f", contr.treatment)"

        return caller.any_wrapper(
            df=df, columns=columns, clause=clause, prefix=prefix, suffix=suffix
        )

    @property
    def columns(self):
        """Get all variables required by the formula."""
        return self.columns_from_formula()

    def columns_from_formula(self=None, formula: str = "") -> list[str]:
        """
        Extract variable names from a formula.

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".

        Returns
        -------
        list[str]
            Variable names required by the formula.
        """
        if type(self) is str:
            formula = self
            self = None

        if self is not None and formula == "":
            formula = self.formula
        return list(Formula(formula).required_variables)

    def lhs(self=None, formula: str = "") -> str:
        """
        Get left-hand side of formula.

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".

        Returns
        -------
        str
            Left-hand side (response variable).
        """
        if type(self) is str:
            formula = self
            self = None

        if self is not None:
            if formula == "":
                formula = self.formula

        #   Separate into subclauses
        sides = formula.split("~")

        if len(sides) == 2:
            lhs_string = sides[0]
        else:
            lhs_string = ""

        return lhs_string

    def rhs(self=None, formula: str = "") -> str:
        """
        Get right-hand side of formula.

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".

        Returns
        -------
        str
            Right-hand side (predictors).
        """
        if type(self) is str:
            formula = self
            self = None

        if self is not None:
            if formula == "":
                formula = self.formula

        #   Separate into subclauses
        sides = formula.split("~")

        if len(sides) == 2:
            rhs_string = sides[1]
        else:
            rhs_string = sides[0]

        return rhs_string

    @property
    def columns_rhs(self):
        """Variables in right-hand side, ordered as in dataframe."""
        formula_rhs = self.rhs()
        columns = FormulaBuilder.columns_from_formula(formula=formula_rhs)

        if self.df is not None:
            return _columns_original_order(
                columns_unordered=columns,
                columns_ordered=nw.from_native(self.df).lazy().collect_schema().names(),
            )
        else:
            return columns

    @property
    def columns_lhs(self):
        """Variables in left-hand side, ordered as in dataframe."""
        formula_lhs = self.lhs()
        if formula_lhs == "":
            return []
        else:
            columns = FormulaBuilder.columns_from_formula(formula=formula_lhs)
            if len(columns) <= 1:
                return columns

            if self.df is not None:
                return _columns_original_order(
                    columns_unordered=columns,
                    columns_ordered=nw.from_native(self.df)
                    .lazy()
                    .collect_schema()
                    .names(),
                )
            else:
                return columns

    def has_constant(
        self=None, formula: str = "", true_if_missing: bool = False
    ) -> bool:
        """
        Check if formula includes an intercept.

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".
        true_if_missing : bool, optional
            Return True if constant term is implicit (not specified).
            Default is False.

        Returns
        -------
        bool
            True if formula includes intercept.
        """
        if type(self) is str:
            formula = self
            self = None

        if self is not None:
            if formula == "":
                formula = self.formula

        #   Separate into subclauses
        sides = formula.split("~")
        if len(sides) == 2:
            lhs = sides[0]
            rhs = sides[1]
        else:
            lhs = None
            rhs = sides[0]

        if true_if_missing:
            return not rhs.replace(" ", "").startswith("0")
        else:
            return rhs.replace(" ", "").startswith("1")

    def remove_constant(self):
        """Remove intercept from formula (change ~1+ to ~0+)."""
        self.formula = self.formula.replace("~1+", "~")
        self.formula = self.formula.replace("~0+", "~")
        self.formula = self.formula.replace("~", "~0+")

    def expand(self=None, formula: str = ""):
        """
        Expand formula shorthand (e.g., a*b becomes a+b+a:b).

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".

        Returns
        -------
        str
            Expanded formula string.
        """
        if self is not None:
            if formula == "":
                formula = self.formula
        else:
            self = FormulaBuilder(formula=formula)

        lhs = self.lhs()
        rhs = self.rhs()

        parser = DefaultFormulaParser()
        parsed = parser.get_terms(rhs)
        reconstructed_formula = "+".join([str(i) for i in parsed])

        if lhs != "":
            reconstructed_formula = f"{lhs}~{reconstructed_formula}"

        if self is not None:
            self.formula = reconstructed_formula

        return reconstructed_formula

    def exclude_interactions(
        self=None,
        formula: str = "",
        b_exclude_powers: bool = True,
        df: IntoFrameT | None = None,
    ) -> tuple[str, bool]:
        """
        Remove interaction terms from formula.

        Parameters
        ----------
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".
        b_exclude_powers : bool, optional
            Also exclude polynomial terms. Default is True.
        df : IntoFrameT | None, optional
            Reference dataframe. Default is None.

        Returns
        -------
        tuple[str, bool]
            (modified formula, whether any terms were dropped)
        """
        if self is not None:
            if df is None:
                df = self.df
            if formula == "":
                formula = self.formula
        else:
            self = FormulaBuilder(df=df, formula=formula)

        #   It's easier with the expanded formula
        self.expand()

        #   Separate into subclauses
        sides = self.formula.split("~")
        if len(sides) == 2:
            lhs = sides[0]
            rhs = sides[1]
        else:
            lhs = ""
            rhs = sides[0]
        subclauses = rhs.split("+")

        rhs = ""

        any_dropped = False

        for clausei in subclauses:
            if b_exclude_powers and "^" in clausei:
                any_dropped = True
                logger.info(
                    f"Dropping {clausei} from formula for having a variable to a power"
                )
            elif ":" in clausei:
                #   Direct interaction
                any_dropped = True
                logger.info(f"Dropping {clausei} from formula")
            else:
                #   include this in the final formula
                rhs += f"+{clausei.strip()}"

        #   Get rid of leading +
        rhs = rhs[1 : len(rhs)]

        output = f"{lhs}~{rhs}"

        self.formula = output
        return (output, any_dropped)

    def exclude_variables(
        self=None,
        exclude_list: list = None,
        formula: str = "",
        df: IntoFrameT | None = None,
        case_insensitive: bool = False,
    ):
        """
        Remove specific variables from formula.

        Parameters
        ----------
        exclude_list : list, optional
            Variables to exclude. Default is None.
        formula : str, optional
            Formula string. If empty, uses self.formula. Default is "".
        df : IntoFrameT | None, optional
            Reference dataframe. Default is None.
        case_insensitive : bool, optional
            Case-insensitive matching. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Modified formula string or self for chaining.
        """
        if exclude_list is None:
            exclude_list = []

        if self is not None:
            if df is None:
                df = self.df
            if formula == "":
                formula = self.formula

        if df is not None:
            exclude_list = columns_from_list(
                df=df, columns=exclude_list, case_insensitive=case_insensitive
            )

        #   Separate into subclauses
        sides = formula.split("~")
        if len(sides) == 2:
            lhs = sides[0]
            rhs = sides[1]
        else:
            lhs = None
            rhs = sides[0]
        subclauses = rhs.split("+")

        rhs = ""

        regex_list = [
            f"(^|([^a-zA-Z0-9_])){itemi}($|([^a-zA-Z0-9_]))" for itemi in exclude_list
        ]
        regexes = "(" + ")|(".join(regex_list) + ")"

        for clausei in subclauses:
            if re.match(regexes, clausei) is None:
                #   include this in the final formula
                rhs += f"+{clausei}"
            else:
                logger.info(f"Dropping {clausei} from formula")

        #   Get rid of leading +
        rhs = rhs[1 : len(rhs)]

        if lhs is not None:
            output = f"{lhs}~{rhs}"
        else:
            output = rhs

        if self is None:
            return output
        else:
            self.formula = output
            return self

    def formula_with_varnames_in_brackets(
        self=None,
        clause: str = "",
        df: pl.LazyFrame | pl.DataFrame | None = None,
        case_insensitive: bool = False,
        append: bool = False,
    ) -> str | FormulaBuilder:
        """
        Expand {pattern} placeholders with matching column names.

        Replaces {var*} with all columns matching var* pattern in the dataframe.
        Useful for programmatically building formulas with wildcards.

        Parameters
        ----------
        clause : str, optional
            Formula clause with {pattern} placeholders. Default is "".
        df : pl.LazyFrame | pl.DataFrame | None, optional
            Dataframe for column lookup. Default is None.
        case_insensitive : bool, optional
            Case-insensitive pattern matching. Default is False.
        append : bool, optional
            Append to existing formula instead of replacing.
            Only used when called on instance. Default is False.

        Returns
        -------
        str | FormulaBuilder
            Expanded formula string or self for chaining.

        Examples
        --------
        >>> fb = FormulaBuilder(df=df)
        >>> fb.formula_with_varnames_in_brackets("{income_*} + age")
        >>> print(fb.formula)
        'y~1+income_wages+income_self_employment+age'
        """

        call_recursively = False
        if df is None and self is not None:
            df = self.df

        #   Separate into subclauses
        sides = clause.split("~")
        if len(sides) == 2:
            lhs = sides[0]
            rhs = sides[1]
        else:
            lhs = None
            rhs = sides[0]
        subclauses = rhs.split("+")

        rhs = ""
        for clausei in subclauses:
            replaced_clause = ""

            left_bracket = clausei.find("{")
            right_bracket = clausei.find("}")

            if left_bracket >= 0 and right_bracket >= 0:
                replace_string = clausei[left_bracket : right_bracket + 1]
                var_name = replace_string[1 : len(replace_string) - 1]
                Columns = columns_from_list(
                    df=df, columns=[var_name], case_insensitive=case_insensitive
                )

                for coli in Columns:
                    if replaced_clause != "":
                        replaced_clause += "+"
                    replaced_clause += clausei.replace(replace_string, coli)

                clausei = replaced_clause

            call_recursively = clausei.find("{") >= 0 and clausei.find("}") >= 0
            rhs += f"+{clausei}"

        #   Get rid of leading +
        rhs = rhs[1 : len(rhs)]

        if lhs is not None:
            output = f"{lhs}~{rhs}"
        else:
            output = rhs

        if call_recursively:
            output = FormulaBuilder.formula_with_varnames_in_brackets(
                clause=output, df=df, case_insensitive=case_insensitive
            )

        if self is None:
            return output
        else:
            if append:
                self.add_to_formula(output)
            else:
                self.formula = output
            return self

columns property

columns

Get all variables required by the formula.

columns_lhs property

columns_lhs

Variables in left-hand side, ordered as in dataframe.

columns_rhs property

columns_rhs

Variables in right-hand side, ordered as in dataframe.

add_to_formula

add_to_formula(
    add_part: str = "", plus_first: bool = True
) -> None

Append a string to the formula.

Parameters:

Name	Type	Description	Default
add_part	str	String to append. Default is "".	''
plus_first	bool	Add "+" before the string. Default is True.	True

Source code in src\survey_kit\utilities\formula_builder.py

def add_to_formula(self, add_part: str = "", plus_first: bool = True) -> None:
    """
    Append a string to the formula.

    Parameters
    ----------
    add_part : str, optional
        String to append. Default is "".
    plus_first : bool, optional
        Add "+" before the string. Default is True.
    """
    if plus_first:
        plus = "+"
    else:
        plus = ""
    self.formula += f"{plus}{add_part}"

any_wrapper

any_wrapper(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
    prefix: str = "",
    suffix: str = "",
) -> str | FormulaBuilder

General wrapper for adding columns with prefix/suffix.

Used internally by other methods to add variables with transformations.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for resolving column patterns. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause (bypasses column lookup). Default is "".	''
case_insensitive	bool	Case-insensitive column matching. Default is False.	False
prefix	str	String to prepend to each column. Default is "".	''
suffix	str	String to append to each column. Default is "".	''

Returns:

Type	Description
str | FormulaBuilder	Formula string if self is None, otherwise returns self for chaining.

Source code in src\survey_kit\utilities\formula_builder.py

def any_wrapper(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
    prefix: str = "",
    suffix: str = "",
) -> str | FormulaBuilder:
    """
    General wrapper for adding columns with prefix/suffix.

    Used internally by other methods to add variables with transformations.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for resolving column patterns. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause (bypasses column lookup). Default is "".
    case_insensitive : bool, optional
        Case-insensitive column matching. Default is False.
    prefix : str, optional
        String to prepend to each column. Default is "".
    suffix : str, optional
        String to append to each column. Default is "".

    Returns
    -------
    str | FormulaBuilder
        Formula string if self is None, otherwise returns self for chaining.
    """
    columns = list_input(columns)

    #   Dataframe to look for columns in
    if df is None:
        df = self.df

    if clause != "":
        output = f"+{prefix}{clause}{suffix}"
    else:
        if df is not None:
            columns = columns_from_list(
                df=df, columns=columns, case_insensitive=case_insensitive
            )

        out_list = [f"{prefix}{coli}{suffix}" for coli in columns]
        output = "+" + "+".join(out_list)

    if self is None:
        return output
    else:
        self.formula += output
        return self

center

center(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder

Center variables (subtract mean).

Convenience method for scale(standardize=False).

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Source code in src\survey_kit\utilities\formula_builder.py

def center(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder:
    """
    Center variables (subtract mean).

    Convenience method for scale(standardize=False).

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    return caller.scale(
        df=df,
        columns=columns,
        clause=clause,
        case_insensitive=case_insensitive,
        standardize=False,
    )

columns_from_formula

columns_from_formula(formula: str = '') -> list[str]

Extract variable names from a formula.

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''

Returns:

Type	Description
list[str]	Variable names required by the formula.

Source code in src\survey_kit\utilities\formula_builder.py

def columns_from_formula(self=None, formula: str = "") -> list[str]:
    """
    Extract variable names from a formula.

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".

    Returns
    -------
    list[str]
        Variable names required by the formula.
    """
    if type(self) is str:
        formula = self
        self = None

    if self is not None and formula == "":
        formula = self.formula
    return list(Formula(formula).required_variables)

continuous

continuous(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder

Add continuous variables to the formula.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns (e.g., "income_*"). Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb = FormulaBuilder(df=df, lhs="y")
>>> fb.continuous(columns=["age", "income"])
>>> print(fb.formula)
'y~1+age+income'

Source code in src\survey_kit\utilities\formula_builder.py

def continuous(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder:
    """
    Add continuous variables to the formula.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns (e.g., "income_*"). Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb = FormulaBuilder(df=df, lhs="y")
    >>> fb.continuous(columns=["age", "income"])
    >>> print(fb.formula)
    'y~1+age+income'
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    return caller.any_wrapper(
        df=df, columns=columns, clause=clause, case_insensitive=case_insensitive
    )

exclude_interactions

exclude_interactions(
    formula: str = "",
    b_exclude_powers: bool = True,
    df: IntoFrameT | None = None,
) -> tuple[str, bool]

Remove interaction terms from formula.

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''
b_exclude_powers	bool	Also exclude polynomial terms. Default is True.	True
df	IntoFrameT | None	Reference dataframe. Default is None.	None

Returns:

Type	Description
tuple[str, bool]	(modified formula, whether any terms were dropped)

Source code in src\survey_kit\utilities\formula_builder.py

def exclude_interactions(
    self=None,
    formula: str = "",
    b_exclude_powers: bool = True,
    df: IntoFrameT | None = None,
) -> tuple[str, bool]:
    """
    Remove interaction terms from formula.

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".
    b_exclude_powers : bool, optional
        Also exclude polynomial terms. Default is True.
    df : IntoFrameT | None, optional
        Reference dataframe. Default is None.

    Returns
    -------
    tuple[str, bool]
        (modified formula, whether any terms were dropped)
    """
    if self is not None:
        if df is None:
            df = self.df
        if formula == "":
            formula = self.formula
    else:
        self = FormulaBuilder(df=df, formula=formula)

    #   It's easier with the expanded formula
    self.expand()

    #   Separate into subclauses
    sides = self.formula.split("~")
    if len(sides) == 2:
        lhs = sides[0]
        rhs = sides[1]
    else:
        lhs = ""
        rhs = sides[0]
    subclauses = rhs.split("+")

    rhs = ""

    any_dropped = False

    for clausei in subclauses:
        if b_exclude_powers and "^" in clausei:
            any_dropped = True
            logger.info(
                f"Dropping {clausei} from formula for having a variable to a power"
            )
        elif ":" in clausei:
            #   Direct interaction
            any_dropped = True
            logger.info(f"Dropping {clausei} from formula")
        else:
            #   include this in the final formula
            rhs += f"+{clausei.strip()}"

    #   Get rid of leading +
    rhs = rhs[1 : len(rhs)]

    output = f"{lhs}~{rhs}"

    self.formula = output
    return (output, any_dropped)

exclude_variables

exclude_variables(
    exclude_list: list = None,
    formula: str = "",
    df: IntoFrameT | None = None,
    case_insensitive: bool = False,
)

Remove specific variables from formula.

Parameters:

Name	Type	Description	Default
exclude_list	list	Variables to exclude. Default is None.	None
formula	str	Formula string. If empty, uses self.formula. Default is "".	''
df	IntoFrameT | None	Reference dataframe. Default is None.	None
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Modified formula string or self for chaining.

Source code in src\survey_kit\utilities\formula_builder.py

def exclude_variables(
    self=None,
    exclude_list: list = None,
    formula: str = "",
    df: IntoFrameT | None = None,
    case_insensitive: bool = False,
):
    """
    Remove specific variables from formula.

    Parameters
    ----------
    exclude_list : list, optional
        Variables to exclude. Default is None.
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".
    df : IntoFrameT | None, optional
        Reference dataframe. Default is None.
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Modified formula string or self for chaining.
    """
    if exclude_list is None:
        exclude_list = []

    if self is not None:
        if df is None:
            df = self.df
        if formula == "":
            formula = self.formula

    if df is not None:
        exclude_list = columns_from_list(
            df=df, columns=exclude_list, case_insensitive=case_insensitive
        )

    #   Separate into subclauses
    sides = formula.split("~")
    if len(sides) == 2:
        lhs = sides[0]
        rhs = sides[1]
    else:
        lhs = None
        rhs = sides[0]
    subclauses = rhs.split("+")

    rhs = ""

    regex_list = [
        f"(^|([^a-zA-Z0-9_])){itemi}($|([^a-zA-Z0-9_]))" for itemi in exclude_list
    ]
    regexes = "(" + ")|(".join(regex_list) + ")"

    for clausei in subclauses:
        if re.match(regexes, clausei) is None:
            #   include this in the final formula
            rhs += f"+{clausei}"
        else:
            logger.info(f"Dropping {clausei} from formula")

    #   Get rid of leading +
    rhs = rhs[1 : len(rhs)]

    if lhs is not None:
        output = f"{lhs}~{rhs}"
    else:
        output = rhs

    if self is None:
        return output
    else:
        self.formula = output
        return self

expand

expand(formula: str = '')

Expand formula shorthand (e.g., a*b becomes a+b+a:b).

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''

Returns:

Type	Description
str	Expanded formula string.

Source code in src\survey_kit\utilities\formula_builder.py

def expand(self=None, formula: str = ""):
    """
    Expand formula shorthand (e.g., a*b becomes a+b+a:b).

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".

    Returns
    -------
    str
        Expanded formula string.
    """
    if self is not None:
        if formula == "":
            formula = self.formula
    else:
        self = FormulaBuilder(formula=formula)

    lhs = self.lhs()
    rhs = self.rhs()

    parser = DefaultFormulaParser()
    parsed = parser.get_terms(rhs)
    reconstructed_formula = "+".join([str(i) for i in parsed])

    if lhs != "":
        reconstructed_formula = f"{lhs}~{reconstructed_formula}"

    if self is not None:
        self.formula = reconstructed_formula

    return reconstructed_formula

factor

factor(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    reference=None,
    case_insensitive: bool = False,
)

Add categorical variables with treatment coding.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
reference	str | int | None	Reference level for treatment coding. Default is None (use first level).	None
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.factor(columns="region")
>>> print(fb.formula)
'y~1+C(region)'

>>> fb.factor(columns="education", reference="high_school")
>>> print(fb.formula)
"y~1+C(education, contr.treatment('high_school'))"

Source code in src\survey_kit\utilities\formula_builder.py

def factor(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    reference=None,
    case_insensitive: bool = False,
):
    """
    Add categorical variables with treatment coding.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    reference : str | int | None, optional
        Reference level for treatment coding. Default is None (use first level).
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.factor(columns="region")
    >>> print(fb.formula)
    'y~1+C(region)'

    >>> fb.factor(columns="education", reference="high_school")
    >>> print(fb.formula)
    "y~1+C(education, contr.treatment('high_school'))"
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    prefix = "C("
    if reference is not None:
        if type(reference) is str:
            suffix = f", contr.treatment('{reference}'))"
        else:
            suffix = f", contr.treatment({reference}))"

    else:
        suffix = ")"  #  f", contr.treatment)"

    return caller.any_wrapper(
        df=df, columns=columns, clause=clause, prefix=prefix, suffix=suffix
    )

formula_with_varnames_in_brackets

formula_with_varnames_in_brackets(
    clause: str = "",
    df: LazyFrame | DataFrame | None = None,
    case_insensitive: bool = False,
    append: bool = False,
) -> str | FormulaBuilder

Expand {pattern} placeholders with matching column names.

Replaces {var} with all columns matching var pattern in the dataframe. Useful for programmatically building formulas with wildcards.

Parameters:

Name	Type	Description	Default
clause	str	Formula clause with {pattern} placeholders. Default is "".	''
df	LazyFrame | DataFrame | None	Dataframe for column lookup. Default is None.	None
case_insensitive	bool	Case-insensitive pattern matching. Default is False.	False
append	bool	Append to existing formula instead of replacing. Only used when called on instance. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Expanded formula string or self for chaining.

Examples:

>>> fb = FormulaBuilder(df=df)
>>> fb.formula_with_varnames_in_brackets("{income_*} + age")
>>> print(fb.formula)
'y~1+income_wages+income_self_employment+age'

Source code in src\survey_kit\utilities\formula_builder.py

def formula_with_varnames_in_brackets(
    self=None,
    clause: str = "",
    df: pl.LazyFrame | pl.DataFrame | None = None,
    case_insensitive: bool = False,
    append: bool = False,
) -> str | FormulaBuilder:
    """
    Expand {pattern} placeholders with matching column names.

    Replaces {var*} with all columns matching var* pattern in the dataframe.
    Useful for programmatically building formulas with wildcards.

    Parameters
    ----------
    clause : str, optional
        Formula clause with {pattern} placeholders. Default is "".
    df : pl.LazyFrame | pl.DataFrame | None, optional
        Dataframe for column lookup. Default is None.
    case_insensitive : bool, optional
        Case-insensitive pattern matching. Default is False.
    append : bool, optional
        Append to existing formula instead of replacing.
        Only used when called on instance. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Expanded formula string or self for chaining.

    Examples
    --------
    >>> fb = FormulaBuilder(df=df)
    >>> fb.formula_with_varnames_in_brackets("{income_*} + age")
    >>> print(fb.formula)
    'y~1+income_wages+income_self_employment+age'
    """

    call_recursively = False
    if df is None and self is not None:
        df = self.df

    #   Separate into subclauses
    sides = clause.split("~")
    if len(sides) == 2:
        lhs = sides[0]
        rhs = sides[1]
    else:
        lhs = None
        rhs = sides[0]
    subclauses = rhs.split("+")

    rhs = ""
    for clausei in subclauses:
        replaced_clause = ""

        left_bracket = clausei.find("{")
        right_bracket = clausei.find("}")

        if left_bracket >= 0 and right_bracket >= 0:
            replace_string = clausei[left_bracket : right_bracket + 1]
            var_name = replace_string[1 : len(replace_string) - 1]
            Columns = columns_from_list(
                df=df, columns=[var_name], case_insensitive=case_insensitive
            )

            for coli in Columns:
                if replaced_clause != "":
                    replaced_clause += "+"
                replaced_clause += clausei.replace(replace_string, coli)

            clausei = replaced_clause

        call_recursively = clausei.find("{") >= 0 and clausei.find("}") >= 0
        rhs += f"+{clausei}"

    #   Get rid of leading +
    rhs = rhs[1 : len(rhs)]

    if lhs is not None:
        output = f"{lhs}~{rhs}"
    else:
        output = rhs

    if call_recursively:
        output = FormulaBuilder.formula_with_varnames_in_brackets(
            clause=output, df=df, case_insensitive=case_insensitive
        )

    if self is None:
        return output
    else:
        if append:
            self.add_to_formula(output)
        else:
            self.formula = output
        return self

function

function(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    operator_before: str = "",
    operator_after: str = "",
    function_item: str = "",
    case_insensitive: bool = False,
    **kwargs,
) -> str | FormulaBuilder

Wrap columns in a function call.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
operator_before	str	String before function call. Default is "".	''
operator_after	str	String after function call. Default is "".	''
function_item	str	Function name (e.g., "log", "sqrt"). Default is "".	''
case_insensitive	bool	Case-insensitive matching. Default is False.	False
**kwargs		Additional function arguments.	{}

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.function(columns="income", function_item="log")
>>> print(fb.formula)
'y~1+log(income)'

Source code in src\survey_kit\utilities\formula_builder.py

def function(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    operator_before: str = "",
    operator_after: str = "",
    function_item: str = "",
    case_insensitive: bool = False,
    **kwargs,
) -> str | FormulaBuilder:
    """
    Wrap columns in a function call.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    operator_before : str, optional
        String before function call. Default is "".
    operator_after : str, optional
        String after function call. Default is "".
    function_item : str, optional
        Function name (e.g., "log", "sqrt"). Default is "".
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.
    **kwargs
        Additional function arguments.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.function(columns="income", function_item="log")
    >>> print(fb.formula)
    'y~1+log(income)'
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    if function_item != "":
        operator_before = f"{operator_before}{function_item}("

        operator_after_final = ""
        if len(kwargs):
            for keyi, valuei in kwargs.items():
                operator_after_final += ","
                if type(valuei) is str:
                    operator_after_final += f"{keyi}='{valuei}'"
                else:
                    operator_after_final += f"{keyi}={valuei}"

        operator_after_final += f"{operator_after})"
    else:
        operator_after_final = operator_after

    return caller.any_wrapper(
        df=df,
        columns=columns,
        clause=clause,
        case_insensitive=case_insensitive,
        prefix=f"{{{operator_before}",
        suffix=f"{operator_after_final}}}",
    )

has_constant

has_constant(
    formula: str = "", true_if_missing: bool = False
) -> bool

Check if formula includes an intercept.

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''
true_if_missing	bool	Return True if constant term is implicit (not specified). Default is False.	False

Returns:

Type	Description
bool	True if formula includes intercept.

Source code in src\survey_kit\utilities\formula_builder.py

def has_constant(
    self=None, formula: str = "", true_if_missing: bool = False
) -> bool:
    """
    Check if formula includes an intercept.

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".
    true_if_missing : bool, optional
        Return True if constant term is implicit (not specified).
        Default is False.

    Returns
    -------
    bool
        True if formula includes intercept.
    """
    if type(self) is str:
        formula = self
        self = None

    if self is not None:
        if formula == "":
            formula = self.formula

    #   Separate into subclauses
    sides = formula.split("~")
    if len(sides) == 2:
        lhs = sides[0]
        rhs = sides[1]
    else:
        lhs = None
        rhs = sides[0]

    if true_if_missing:
        return not rhs.replace(" ", "").startswith("0")
    else:
        return rhs.replace(" ", "").startswith("1")

interact_clauses

interact_clauses(
    clause1: str = "",
    clause2: str = "",
    no_base: bool = False,
) -> str | FormulaBuilder

Create interactions between two formula clauses.

Parameters:

Name	Type	Description	Default
clause1	str	First formula clause. Default is "".	''
clause2	str	Second formula clause. Default is "".	''
no_base	bool	Exclude main effects from clauses. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.interact_clauses("age+education", "region")
>>> print(fb.formula)
'y~1+(age+education)*(region)'

Source code in src\survey_kit\utilities\formula_builder.py

def interact_clauses(
    self=None, clause1: str = "", clause2: str = "", no_base: bool = False
) -> str | FormulaBuilder:
    """
    Create interactions between two formula clauses.

    Parameters
    ----------
    clause1 : str, optional
        First formula clause. Default is "".
    clause2 : str, optional
        Second formula clause. Default is "".
    no_base : bool, optional
        Exclude main effects from clauses. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.interact_clauses("age+education", "region")
    >>> print(fb.formula)
    'y~1+(age+education)*(region)'
    """
    if clause1.startswith("+"):
        clause1 = clause1[1:]
    if clause2.startswith("+"):
        clause2 = clause2[1:]

    clause1 = clause1.replace("++", "+")
    clause2 = clause2.replace("++", "+")

    output = f"({clause1})*({clause2})"

    if no_base:
        output += f"-({clause1} + {clause2})"

    if self is None:
        return f"+{output}"
    else:
        self.add_to_formula(output)
        return self

lhs

lhs(formula: str = '') -> str

Get left-hand side of formula.

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''

Returns:

Type	Description
str	Left-hand side (response variable).

Source code in src\survey_kit\utilities\formula_builder.py

def lhs(self=None, formula: str = "") -> str:
    """
    Get left-hand side of formula.

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".

    Returns
    -------
    str
        Left-hand side (response variable).
    """
    if type(self) is str:
        formula = self
        self = None

    if self is not None:
        if formula == "":
            formula = self.formula

    #   Separate into subclauses
    sides = formula.split("~")

    if len(sides) == 2:
        lhs_string = sides[0]
    else:
        lhs_string = ""

    return lhs_string

polynomial

polynomial(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    degree: int = 0,
    case_insensitive: bool = False,
    center: bool = False,
) -> str | FormulaBuilder

Add polynomial terms.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
degree	int	Polynomial degree. Default is 0 (returns continuous).	0
case_insensitive	bool	Case-insensitive matching. Default is False.	False
center	bool	Use orthogonal polynomials (centered). Default is False (raw polynomials).	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.polynomial(columns="age", degree=3)
>>> print(fb.formula)
'y~1+poly(age,degree=3,raw=True)'

>>> fb.polynomial(columns="age", degree=2, center=True)
>>> print(fb.formula)
'y~1+poly(age,degree=2,raw=False)'

Source code in src\survey_kit\utilities\formula_builder.py

def polynomial(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    degree: int = 0,
    case_insensitive: bool = False,
    center: bool = False,
) -> str | FormulaBuilder:
    """
    Add polynomial terms.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    degree : int, optional
        Polynomial degree. Default is 0 (returns continuous).
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.
    center : bool, optional
        Use orthogonal polynomials (centered). Default is False (raw polynomials).

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.polynomial(columns="age", degree=3)
    >>> print(fb.formula)
    'y~1+poly(age,degree=3,raw=True)'

    >>> fb.polynomial(columns="age", degree=2, center=True)
    >>> print(fb.formula)
    'y~1+poly(age,degree=2,raw=False)'
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self
        if df is None:
            df = self.df

    if degree <= 1:
        if center:
            return caller.center(
                df=df,
                columns=columns,
                clause=clause,
                case_insensitive=case_insensitive,
            )
        else:
            return caller.continuous(
                df=df,
                columns=columns,
                clause=clause,
                case_insensitive=case_insensitive,
            )
    else:
        subformula = ""

        for power in range(1, degree + 1):
            operator_before = "poly("
            if center:
                operator_after = f",degree={degree},raw=False)"
            else:
                operator_after = f",degree={degree},raw=True)"

            subformula += FormulaBuilder.function(
                df=df,
                columns=columns,
                clause=clause,
                operator_before=operator_before,
                operator_after=operator_after,
                case_insensitive=case_insensitive,
            )

        if self is None:
            return subformula
        else:
            self.add_to_formula(subformula, False)
            return self

remove_constant

remove_constant()

Remove intercept from formula (change ~1+ to ~0+).

Source code in src\survey_kit\utilities\formula_builder.py

def remove_constant(self):
    """Remove intercept from formula (change ~1+ to ~0+)."""
    self.formula = self.formula.replace("~1+", "~")
    self.formula = self.formula.replace("~0+", "~")
    self.formula = self.formula.replace("~", "~0+")

rhs

rhs(formula: str = '') -> str

Get right-hand side of formula.

Parameters:

Name	Type	Description	Default
formula	str	Formula string. If empty, uses self.formula. Default is "".	''

Returns:

Type	Description
str	Right-hand side (predictors).

Source code in src\survey_kit\utilities\formula_builder.py

def rhs(self=None, formula: str = "") -> str:
    """
    Get right-hand side of formula.

    Parameters
    ----------
    formula : str, optional
        Formula string. If empty, uses self.formula. Default is "".

    Returns
    -------
    str
        Right-hand side (predictors).
    """
    if type(self) is str:
        formula = self
        self = None

    if self is not None:
        if formula == "":
            formula = self.formula

    #   Separate into subclauses
    sides = formula.split("~")

    if len(sides) == 2:
        rhs_string = sides[1]
    else:
        rhs_string = sides[0]

    return rhs_string

scale

scale(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    standardize: bool = True,
    case_insensitive: bool = False,
) -> str | FormulaBuilder

Standardize or center variables.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
standardize	bool	If True, standardize (mean=0, sd=1). If False, only center (mean=0). Default is True.	True
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.scale(columns="age")  # Standardize
>>> fb.scale(columns="income", standardize=False)  # Center only

Source code in src\survey_kit\utilities\formula_builder.py

def scale(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    standardize: bool = True,
    case_insensitive: bool = False,
) -> str | FormulaBuilder:
    """
    Standardize or center variables.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    standardize : bool, optional
        If True, standardize (mean=0, sd=1). If False, only center (mean=0).
        Default is True.
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.scale(columns="age")  # Standardize
    >>> fb.scale(columns="income", standardize=False)  # Center only
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    if standardize:
        function_item = "scale"
    else:
        function_item = "center"

    return caller.any_wrapper(
        df=df,
        columns=columns,
        clause=clause,
        case_insensitive=case_insensitive,
        prefix=f"{function_item}(",
        suffix=")",
    )

simple_interaction

simple_interaction(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    order: int = 2,
    case_insensitive: bool = False,
    sub_function: Callable | None = None,
    no_base: bool = False,
) -> str | FormulaBuilder

Create interactions between variables.

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
order	int	Interaction order. 2 = pairwise, 3 = three-way, etc. Default is 2.	2
case_insensitive	bool	Case-insensitive matching. Default is False.	False
sub_function	Callable | None	Function to apply to columns before interacting. Default is None.	None
no_base	bool	Exclude main effects (only interactions). Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Examples:

>>> fb.simple_interaction(columns=["age", "education"], order=2)
>>> print(fb.formula)
'y~1+(age+education)**2'

>>> fb.simple_interaction(columns=["a", "b", "c"], order=2, no_base=True)
>>> print(fb.formula)
'y~1+(a+b+c)**2-(a+b+c)'  # Interactions only

Source code in src\survey_kit\utilities\formula_builder.py

def simple_interaction(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    order: int = 2,
    case_insensitive: bool = False,
    sub_function: Callable | None = None,
    no_base: bool = False,
) -> str | FormulaBuilder:
    """
    Create interactions between variables.

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    order : int, optional
        Interaction order. 2 = pairwise, 3 = three-way, etc. Default is 2.
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.
    sub_function : Callable | None, optional
        Function to apply to columns before interacting. Default is None.
    no_base : bool, optional
        Exclude main effects (only interactions). Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.

    Examples
    --------
    >>> fb.simple_interaction(columns=["age", "education"], order=2)
    >>> print(fb.formula)
    'y~1+(age+education)**2'

    >>> fb.simple_interaction(columns=["a", "b", "c"], order=2, no_base=True)
    >>> print(fb.formula)
    'y~1+(a+b+c)**2-(a+b+c)'  # Interactions only
    """
    if self is not None:
        if df is None:
            df = self.df

    if sub_function is None:
        sub_function = FormulaBuilder.continuous

    columns = columns_from_list(
        df=df, columns=columns, case_insensitive=case_insensitive
    )

    subformula = sub_function(
        df=df, columns=columns, case_insensitive=case_insensitive
    )
    #   Remove leading plus sign
    subformula = subformula[1 : len(subformula)]

    if len(columns) > 1:
        output = f"({subformula})**{order}"
    else:
        output = subformula

    if no_base:
        output += f"-({subformula})"

    if self is None:
        return f"+{output}"
    else:
        self.add_to_formula(output)
        return self

standardize

standardize(
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder

Standardize variables (mean=0, sd=1).

Convenience method for scale(standardize=True).

Parameters:

Name	Type	Description	Default
df	IntoFrameT | None	Dataframe for column lookup. Default is None.	None
columns	str | list | None	Column names or patterns. Default is None.	None
clause	str	Pre-constructed clause. Default is "".	''
case_insensitive	bool	Case-insensitive matching. Default is False.	False

Returns:

Type	Description
str | FormulaBuilder	Formula string or self for chaining.

Source code in src\survey_kit\utilities\formula_builder.py

def standardize(
    self=None,
    df: IntoFrameT | None = None,
    columns: str | list | None = None,
    clause: str = "",
    case_insensitive: bool = False,
) -> str | FormulaBuilder:
    """
    Standardize variables (mean=0, sd=1).

    Convenience method for scale(standardize=True).

    Parameters
    ----------
    df : IntoFrameT | None, optional
        Dataframe for column lookup. Default is None.
    columns : str | list | None, optional
        Column names or patterns. Default is None.
    clause : str, optional
        Pre-constructed clause. Default is "".
    case_insensitive : bool, optional
        Case-insensitive matching. Default is False.

    Returns
    -------
    str | FormulaBuilder
        Formula string or self for chaining.
    """
    if self is None:
        caller = FormulaBuilder
    else:
        caller = self

    return caller.scale(
        df=df,
        columns=columns,
        clause=clause,
        case_insensitive=case_insensitive,
        standardize=True,
    )