diff --git a/causalml/inference/_arg_order.py b/causalml/inference/_arg_order.py new file mode 100644 index 00000000..476da69b --- /dev/null +++ b/causalml/inference/_arg_order.py @@ -0,0 +1,188 @@ +"""#854: scikit-learn ``fit(X, y, ...)`` argument-order migration. + +CausalML learners historically take ``fit(X, treatment, y, ...)``, which puts +``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final +estimator's ``fit(X, y)`` positionally). In v1.0 the positional order becomes +``(X, y, treatment, ...)``: + +* ``y`` and ``treatment`` move to the front, immediately after ``X``, and every + other parameter keeps its relative position — so + ``BaseDRIVLearner.fit(X, assignment, treatment, y, ...)`` becomes + ``fit(X, y, treatment, assignment, ...)``. +* A suffixed pair such as ``treatment_val``/``y_val`` is reordered in place, so + ``UpliftTreeClassifier.fit(X, treatment, y, X_val, treatment_val, y_val, ...)`` + becomes ``fit(X, y, treatment, X_val, y_val, treatment_val, ...)`` rather than + leaving the validation triple inconsistent with the main one. + +This is a two-step deprecation. Step one (this module): the positional order is +UNCHANGED, so no existing call silently breaks, but passing ``treatment`` or +``y`` positionally emits a ``FutureWarning`` steering callers to keyword +arguments, which are order-independent and therefore safe across the flip. +Step two (v1.0): reorder the signatures and delete this module. + +Coverage is decided by **signature, not by method name**. Any method whose +positional order would change is shimmed — ``fit`` and ``predict``, but equally +``bootstrap``, ``fit_bootstrap_ensemble``, ``prune`` and the ``Sensitivity`` +helpers. An earlier name-based allowlist silently missed eleven public methods +(#981); the deprecation window is one-shot, so a method that is not shimmed here +would need a second deprecation cycle of its own. +""" + +import functools +import inspect +import warnings + +#: Parameters whose positional slot changes in v1.0, in their v1.0 order. +REORDERED_PARAMS = ("y", "treatment") + + +def _positional_params(method): + """Return the names of ``method``'s positional parameters, minus ``self``.""" + kinds = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ) + names = [ + p.name for p in inspect.signature(method).parameters.values() if p.kind in kinds + ] + return names[1:] if names[:1] == ["self"] else names + + +def _order_suffixed_pairs(names): + """Put ``y`` before ``treatment`` within each suffixed pair, in place. + + Only pairs sharing a suffix are touched (``treatment_val``/``y_val``), and + they keep the two slots they already occupy — unlike the bare + ``y``/``treatment`` pair, which is hoisted to the front by :func:`v1_order`. + """ + ordered = list(names) + suffixes = { + name[len("treatment") :] + for name in ordered + if name.startswith("treatment") and name != "treatment" + } + for suffix in suffixes: + pair = ("y" + suffix, "treatment" + suffix) + if all(name in ordered for name in pair): + first, second = sorted(ordered.index(name) for name in pair) + ordered[first], ordered[second] = pair + return ordered + + +def v1_order(params): + """Reorder ``params`` to the v1.0 convention. + + The feature matrix stays first, then ``y`` and ``treatment``, then + everything else in its original relative order with any suffixed + ``treatment``/``y`` pair swapped so ``y`` leads. + + ``params`` is returned unchanged unless **both** ``y`` and ``treatment`` + appear after the first position. This deprecation is about that pair; a + method taking only one of them is not part of it, which keeps the rule off + the vendored sklearn tree builders (``build(tree, X, y, ...)``, where + hoisting ``y`` would put it ahead of ``X``) and off dataset helpers that + name their treatment vector something else. + """ + head, tail = params[:1], params[1:] + if not all(name in tail for name in REORDERED_PARAMS): + return list(params) + moved = list(REORDERED_PARAMS) + rest = [name for name in tail if name not in REORDERED_PARAMS] + return head + moved + _order_suffixed_pairs(rest) + + +def _may_reorder(method): + """Cheap pre-filter: could this method's positional order change at all? + + Reads the code object directly so the common case costs no + :func:`inspect.signature` call — ``shim_arg_order`` scans every method of + every learner class at import time. Unwraps first, because + :func:`functools.wraps` copies ``__wrapped__`` but not ``__code__``, so a + decorated method (e.g. ``@timeit``) would otherwise look argument-less. + """ + code = getattr(inspect.unwrap(method), "__code__", None) + if code is None: + return False + args = set(code.co_varnames[: code.co_argcount]) + return bool(args & set(REORDERED_PARAMS)) + + +def deprecate_positional_treatment_y(method): + """Warn when ``treatment``/``y`` are passed positionally to ``method``. + + ``method`` is returned unchanged when its positional order is already the + v1.0 order (including when it takes neither parameter), so this is safe to + apply blanket-style from ``__init_subclass__``. + + Only the outermost shimmed call on a given instance warns, so internal + delegation (``fit_predict`` -> ``fit``/``predict``, a subclass ``fit`` -> + ``super().fit``, ``estimate_ate`` -> ``fit_predict`` -> ``fit``) never + double-counts and callers see at most one warning per top-level call. + :func:`functools.wraps` preserves the real signature, so introspection and + scikit-learn metadata routing are unaffected. + """ + params = _positional_params(method) + target = v1_order(params) + if target == params: + return method + + message = ( + "Passing `treatment` and/or `y` to {name}() by position is deprecated " + "and will change in causalml v1.0: the positional argument order " + "becomes ({new_order}) for scikit-learn Pipeline compatibility " + "(see https://github.com/uber/causalml/issues/854). To be safe across " + "the change, pass them as keyword arguments, e.g. {example}." + ).format( + name=method.__name__, + new_order=", ".join(target), + example="{}({}, {})".format( + method.__name__, + params[0], + ", ".join(f"{p}={p}" for p in REORDERED_PARAMS if p in params), + ), + ) + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + if getattr(self, "_in_arg_order_call", False): + return method(self, *args, **kwargs) + if any( + params[i] in REORDERED_PARAMS for i in range(min(len(args), len(params))) + ): + warnings.warn(message, FutureWarning, stacklevel=2) + self._in_arg_order_call = True + try: + return method(self, *args, **kwargs) + finally: + self._in_arg_order_call = False + + wrapper._arg_order_shimmed = True + return wrapper + + +def shim_arg_order(cls): + """Shim every method of ``cls`` whose positional order changes in v1.0. + + Membership is decided by signature rather than by a list of method names — + a name list is what let ``bootstrap``, ``prune`` and the ``Sensitivity`` + helpers slip through the window (#981). Only methods defined directly on + ``cls`` are considered (inherited ones are already wrapped on the parent); + dunders, ``classmethod``/``staticmethod`` descriptors, abstract methods and + already-wrapped methods are skipped. + + ``SerializableLearner`` applies this automatically via ``__init_subclass__``. + Classes outside that hierarchy — the TF/Torch/JAX estimators, + ``PolicyLearner``, ``TMLELearner``, ``Sensitivity`` — use it as a decorator. + """ + for name, method in list(vars(cls).items()): + if ( + name.startswith("_") + or isinstance(method, (classmethod, staticmethod)) + or not callable(method) + or getattr(method, "__isabstractmethod__", False) + or getattr(method, "_arg_order_shimmed", False) + or not _may_reorder(method) + ): + continue + setattr(cls, name, deprecate_positional_treatment_y(method)) + return cls diff --git a/causalml/inference/jax/cevae/cevae.py b/causalml/inference/jax/cevae/cevae.py index 43ff8d1c..cb75efb9 100644 --- a/causalml/inference/jax/cevae/cevae.py +++ b/causalml/inference/jax/cevae/cevae.py @@ -24,6 +24,7 @@ import optax from causalml.inference.meta.utils import convert_pd_to_np +from causalml.inference._arg_order import shim_arg_order from causalml.inference.jax.cevae.losses import cevae_loss from causalml.inference.jax.cevae.modeling import Guide, Model, PreWhitener @@ -80,6 +81,7 @@ def _loss(net): return train_step +@shim_arg_order class CEVAE: """JAX/flax.nnx CEVAE for treatment-effect estimation. diff --git a/causalml/inference/jax/dragonnet.py b/causalml/inference/jax/dragonnet.py index 0b53335e..ab4f7607 100644 --- a/causalml/inference/jax/dragonnet.py +++ b/causalml/inference/jax/dragonnet.py @@ -25,6 +25,7 @@ make_tarreg_loss, ) from causalml.inference.meta.utils import convert_pd_to_np +from causalml.inference._arg_order import shim_arg_order class EpsilonLayer(nnx.Module): @@ -278,6 +279,7 @@ def _run_training_loop( break +@shim_arg_order class DragonNet: """JAX/flax.nnx DragonNet for treatment effect estimation. diff --git a/causalml/inference/meta/base.py b/causalml/inference/meta/base.py index c8792166..e50d40cb 100644 --- a/causalml/inference/meta/base.py +++ b/causalml/inference/meta/base.py @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod import logging + import numpy as np import pandas as pd from joblib import Parallel, delayed diff --git a/causalml/inference/meta/tmle.py b/causalml/inference/meta/tmle.py index 372d0f3d..d1363c45 100644 --- a/causalml/inference/meta/tmle.py +++ b/causalml/inference/meta/tmle.py @@ -11,6 +11,7 @@ check_p_conditions, convert_pd_to_np, ) +from causalml.inference._arg_order import shim_arg_order logger = logging.getLogger("causalml") @@ -91,6 +92,7 @@ def simple_tmle(y, w, q0w, q1w, p, alpha=0.0001): return np.mean(q1star - q0star), np.sqrt(np.var(ic) / np.size(y)) +@shim_arg_order class TMLELearner: """Targeted maximum likelihood estimation. diff --git a/causalml/inference/serialization.py b/causalml/inference/serialization.py index cdcf8bc3..1bfc8557 100644 --- a/causalml/inference/serialization.py +++ b/causalml/inference/serialization.py @@ -14,6 +14,8 @@ import joblib +from causalml.inference._arg_order import shim_arg_order + logger = logging.getLogger("causalml") @@ -44,6 +46,17 @@ class SerializableLearner: check does not apply (e.g. for non-sklearn learners). """ + def __init_subclass__(cls, **kwargs): + """Apply the #854 argument-order shim to each subclass. + + Hosting this on the shared mixin rather than on a single learner base + covers the meta-learners, the causal/uplift trees and forests, and the + IV learners in one place. See :func:`~causalml.inference._arg_order. + shim_arg_order` for what gets wrapped. + """ + super().__init_subclass__(**kwargs) + shim_arg_order(cls) + def _is_fitted(self): """Check whether this learner has been fitted. diff --git a/causalml/inference/tf/dragonnet.py b/causalml/inference/tf/dragonnet.py index f8cd3ed3..72ae4c75 100644 --- a/causalml/inference/tf/dragonnet.py +++ b/causalml/inference/tf/dragonnet.py @@ -33,8 +33,10 @@ make_tarreg_loss, ) from causalml.inference.meta.utils import convert_pd_to_np +from causalml.inference._arg_order import shim_arg_order +@shim_arg_order class DragonNet: def __init__( self, diff --git a/causalml/inference/torch/cevae.py b/causalml/inference/torch/cevae.py index 3e7b0f79..68d6e488 100644 --- a/causalml/inference/torch/cevae.py +++ b/causalml/inference/torch/cevae.py @@ -27,6 +27,7 @@ from pyro.contrib.cevae import CEVAE as CEVAEModel from causalml.inference.meta.utils import convert_pd_to_np +from causalml.inference._arg_order import shim_arg_order pyro_logger = logging.getLogger("pyro") pyro_logger.setLevel(logging.DEBUG) @@ -34,6 +35,7 @@ pyro_logger.handlers[0].setLevel(logging.DEBUG) +@shim_arg_order class CEVAE: def __init__( self, diff --git a/causalml/inference/tree/utils.py b/causalml/inference/tree/utils.py index fbd22efe..0baa09d8 100644 --- a/causalml/inference/tree/utils.py +++ b/causalml/inference/tree/utils.py @@ -2,6 +2,7 @@ Utility functions for uplift trees. """ +import functools import time from typing import Callable @@ -342,6 +343,7 @@ def timeit(exclude_kwargs: tuple = ()) -> Callable: """ def wrapper(f: Callable): + @functools.wraps(f) def wrapped(*args, **kw): ts = time.time() result = f(*args, **kw) diff --git a/causalml/metrics/sensitivity.py b/causalml/metrics/sensitivity.py index c0ac4206..2ef0830e 100644 --- a/causalml/metrics/sensitivity.py +++ b/causalml/metrics/sensitivity.py @@ -4,6 +4,8 @@ import matplotlib.pyplot as plt from importlib import import_module +from causalml.inference._arg_order import shim_arg_order + logger = logging.getLogger("sensitivity") SUMMARY_COLS = ["Method", "ATE", "New ATE", "New ATE LB", "New ATE UB"] @@ -95,6 +97,7 @@ def msm_propensity_bounds(p, gamma): return p_lower, p_upper +@shim_arg_order class Sensitivity: """A Sensitivity Check class to support Placebo Treatment, Irrelevant Additional Confounder and Subset validation refutation methods to verify causal inference. diff --git a/causalml/optimize/policylearner.py b/causalml/optimize/policylearner.py index 1dea008a..349cc669 100644 --- a/causalml/optimize/policylearner.py +++ b/causalml/optimize/policylearner.py @@ -2,6 +2,7 @@ import numpy as np from causalml.propensity import compute_propensity_score +from causalml.inference._arg_order import shim_arg_order from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier from sklearn.model_selection import KFold from sklearn.tree import DecisionTreeClassifier @@ -9,6 +10,7 @@ logger = logging.getLogger("causalml") +@shim_arg_order class PolicyLearner: """ A Learner that learns a treatment assignment policy with observational data using doubly robust estimator of causal diff --git a/tests/test_fit_arg_order.py b/tests/test_fit_arg_order.py new file mode 100644 index 00000000..c7ea4990 --- /dev/null +++ b/tests/test_fit_arg_order.py @@ -0,0 +1,381 @@ +"""Tests for the #854 argument-order deprecation shim. + +CausalML learners historically take ``fit(X, treatment, y, ...)``, which puts +``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final +estimator's ``fit(X, y)`` positionally). For scikit-learn compatibility the +positional order becomes ``(X, y, treatment, ...)`` in v1.0 — across the whole +package (meta-learners, causal/uplift trees and forests, IV, and the TF/Torch/ +JAX estimators), and for ``predict`` as well as the fit family, so the same two +arguments never sit in opposite orders on one class. + +This is a two-step deprecation: the positional order is UNCHANGED for now (so no +existing call silently breaks), but passing ``treatment``/``y`` positionally emits +a ``FutureWarning`` steering callers to keyword arguments, which are +order-independent and therefore safe across the v1.0 flip. These tests pin that +behavior; they should be updated (not deleted) when the signatures are reordered. +""" + +import importlib +import inspect +import pkgutil +import warnings + +import numpy as np +import pytest +from sklearn.linear_model import LinearRegression + +from causalml.inference._arg_order import _positional_params, v1_order +from causalml.inference.iv import BaseDRIVRegressor, IVRegressor +from causalml.inference.meta import ( + BaseSRegressor, + BaseTRegressor, + BaseXRegressor, + BaseRRegressor, + BaseDRRegressor, + XGBTRegressor, + BaseSClassifier, + BaseTClassifier, + BaseXClassifier, + BaseRClassifier, +) +from causalml.inference.tree import CausalTreeRegressor, UpliftTreeClassifier +from causalml.metrics.sensitivity import Sensitivity + +from .const import RANDOM_SEED + +# One representative regressor per meta-learner family; all default to +# ``control_name=0`` which matches the 0/1 treatment from ``synthetic_data``. +REGRESSORS = [ + BaseSRegressor, + BaseTRegressor, + BaseXRegressor, + BaseRRegressor, + BaseDRRegressor, +] + +# The R-learner's ``predict(X, p, return_components)`` already omits treatment/y, +# so it is unaffected by the flip. Derive the list instead of hard-coding it, so +# it stays correct if a predict signature changes. +PREDICT_REORDERED = [ + cls + for cls in REGRESSORS + if "treatment" + in _positional_params(getattr(cls.predict, "__wrapped__", cls.predict)) +] + + +def _order_warnings(record): + """Filter a warning record down to the #854 arg-order FutureWarnings.""" + return [ + w + for w in record + if issubclass(w.category, FutureWarning) and "argument order" in str(w.message) + ] + + +@pytest.mark.parametrize("learner_cls", REGRESSORS) +def test_positional_fit_warns(generate_regression_data, learner_cls): + y, X, treatment, _, _, _ = generate_regression_data() + learner = learner_cls(learner=LinearRegression()) + with pytest.warns(FutureWarning, match="argument order"): + learner.fit(X, treatment, y) + + +@pytest.mark.parametrize("learner_cls", REGRESSORS) +def test_keyword_fit_does_not_warn(generate_regression_data, learner_cls): + y, X, treatment, _, _, _ = generate_regression_data() + learner = learner_cls(learner=LinearRegression()) + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.fit(X=X, y=y, treatment=treatment) + assert _order_warnings(record) == [] + + +def test_positional_and_keyword_are_equivalent(generate_regression_data): + """The shim must never silently swap y/treatment: both call styles must fit + identically. Guards against a values-based reorder that could corrupt data.""" + y, X, treatment, _, _, _ = generate_regression_data() + + positional = BaseTRegressor(learner=LinearRegression()) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + positional.fit(X, treatment, y) + + keyword = BaseTRegressor(learner=LinearRegression()) + keyword.fit(X=X, y=y, treatment=treatment) + + np.testing.assert_allclose(positional.predict(X), keyword.predict(X)) + + +def test_fit_predict_warns_once(generate_regression_data): + """fit_predict -> fit delegation must not double-count the warning.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = BaseTRegressor(learner=LinearRegression()) + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.fit_predict(X, treatment, y) + assert len(_order_warnings(record)) == 1 + + +def test_estimate_ate_warns_once(generate_regression_data): + """estimate_ate -> (fit_predict ->) fit delegation must warn exactly once.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = BaseSRegressor(learner=LinearRegression()) + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.estimate_ate(X, treatment, y) + assert len(_order_warnings(record)) == 1 + + +def test_subclass_super_fit_warns_once(generate_regression_data): + """A subclass fit that delegates to super().fit (XGBTRegressor) warns once.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = XGBTRegressor() + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.fit(X, treatment, y) + assert len(_order_warnings(record)) == 1 + + +@pytest.mark.parametrize("learner_cls", REGRESSORS) +def test_fit_signature_is_preserved(learner_cls): + """functools.wraps keeps the real (X, treatment, y, ...) signature visible.""" + learner = learner_cls(learner=LinearRegression()) + params = list(inspect.signature(learner.fit).parameters) + assert params[:3] == ["X", "treatment", "y"] + + +# --- v1.0 target ordering --------------------------------------------------- + + +@pytest.mark.parametrize( + "current, expected", + [ + # Meta-learners: treatment and y simply swap. + (["X", "treatment", "y", "p"], ["X", "y", "treatment", "p"]), + # predict: same rule, so fit and predict agree after the flip. + ( + ["X", "treatment", "y", "p", "return_components", "verbose"], + ["X", "y", "treatment", "p", "return_components", "verbose"], + ), + # IV: the instrument keeps its relative position and shifts right. + (["X", "treatment", "y", "w"], ["X", "y", "treatment", "w"]), + # DRIV: `assignment` lands fourth, after X, y and treatment. + ( + ["X", "assignment", "treatment", "y", "p", "pZ"], + ["X", "y", "treatment", "assignment", "p", "pZ"], + ), + # UpliftTreeClassifier: the validation triple is reordered in place too, + # so it stays consistent with the main one (#980 open question 1). + ( + ["X", "treatment", "y", "X_val", "treatment_val", "y_val", "sample_weight"], + ["X", "y", "treatment", "X_val", "y_val", "treatment_val", "sample_weight"], + ), + # Sensitivity helpers are in scope, and `p` shifts right (#980 q2). + (["X", "p", "treatment", "y"], ["X", "y", "treatment", "p"]), + # A signature with neither is left completely alone. + (["X", "sample_weight"], ["X", "sample_weight"]), + # Only `y`, no `treatment`: not this deprecation, so untouched. Keeps the + # rule off the vendored sklearn tree builders, where hoisting `y` would + # put it ahead of `X`. + (["tree", "X", "y", "sample_weight"], ["tree", "X", "y", "sample_weight"]), + (["X", "w", "y", "initial_taus"], ["X", "w", "y", "initial_taus"]), + ], +) +def test_v1_order(current, expected): + """X first, then y, then treatment; everything else keeps its relative order.""" + assert v1_order(current) == expected + + +# --- coverage across the whole package -------------------------------------- + + +def _shimmable_methods(): + """Yield (qualname, method) for every public method whose order changes at v1.0. + + Scans by **signature**, not by method name. An earlier name-based allowlist + silently missed eleven public methods (#981) — `bootstrap`, + `fit_bootstrap_ensemble`, `bootstrap_pool`, `prune`, `fill` and the three + `Sensitivity.get_*` helpers — so this guard must not be narrowed back to a + fixed set of names. + + Walks the installed package so a newly added learner is picked up + automatically. Modules whose optional backend (tf/torch/jax) is missing are + skipped, so this covers them only in the backend CI lanes. + """ + import causalml + + for info in pkgutil.walk_packages(causalml.__path__, prefix="causalml."): + try: + module = importlib.import_module(info.name) + except Exception: # optional backend absent, or an import-time failure + continue + for _, cls in inspect.getmembers(module, inspect.isclass): + if not cls.__module__.startswith("causalml."): + continue + for name, method in vars(cls).items(): + if ( + name.startswith("_") + or isinstance(method, (classmethod, staticmethod)) + or not callable(method) + or getattr(method, "__isabstractmethod__", False) + ): + continue + try: + params = _positional_params(getattr(method, "__wrapped__", method)) + except (TypeError, ValueError): # not introspectable + continue + if v1_order(params) != params: + yield f"{cls.__module__}.{cls.__qualname__}.{name}", method + + +def test_every_reordered_method_is_shimmed(): + """Completeness guard for the deprecation window. + + A deprecation window is one-shot: anything whose positional order changes at + v1.0 must warn in this release or it needs a second cycle. This fails if any + public method's order would change without the shim. + """ + unshimmed = [ + qualname + for qualname, method in _shimmable_methods() + if not getattr(method, "_arg_order_shimmed", False) + ] + assert unshimmed == [], f"missing #854 shim on: {unshimmed}" + + +def test_shim_reaches_past_the_fit_family(): + """#981: the shim must not regress to a fixed list of method names.""" + names = {qualname.rsplit(".", 1)[1] for qualname, _ in _shimmable_methods()} + assert {"bootstrap", "fit_bootstrap_ensemble", "get_prediction"} <= names, names + + +def test_shim_covers_more_than_the_meta_learners(): + """Guards against the hook silently regressing to BaseLearner-only scope.""" + modules = {qualname.rsplit(".", 2)[0] for qualname, _ in _shimmable_methods()} + assert any("inference.tree" in m for m in modules), modules + assert any("inference.iv" in m for m in modules), modules + + +# --- behaviour outside the meta-learner family ------------------------------ + + +def test_causal_tree_fit_warns_positionally(generate_regression_data): + """The trees inherit sklearn estimator bases, so they need the flip too.""" + y, X, treatment, _, _, _ = generate_regression_data() + with pytest.warns(FutureWarning, match="argument order"): + CausalTreeRegressor().fit(X, treatment, y) + + +def test_causal_tree_fit_keyword_is_silent(generate_regression_data): + y, X, treatment, _, _, _ = generate_regression_data() + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + CausalTreeRegressor().fit(X=X, y=y, treatment=treatment) + assert _order_warnings(record) == [] + + +def test_iv_regressor_fit_warns_positionally(): + """IVRegressor's message must name its own signature, instrument included.""" + rng = np.random.RandomState(RANDOM_SEED) + n = 200 + w = rng.binomial(1, 0.5, n).astype(float) + treatment = w * rng.binomial(1, 0.8, n) + X = rng.normal(size=(n, 2)) + y = treatment + X[:, 0] + rng.normal(size=n) + with pytest.warns(FutureWarning, match=r"becomes \(X, y, treatment, w\)"): + IVRegressor().fit(X, treatment, y, w) + + +def test_driv_message_puts_assignment_fourth(): + """Pins the v1.0 order chosen for DRIV: (X, y, treatment, assignment, ...).""" + params = _positional_params(BaseDRIVRegressor.fit.__wrapped__) + assert v1_order(params)[:4] == ["X", "y", "treatment", "assignment"] + + +# --- predict joins the deprecation ------------------------------------------ + + +@pytest.mark.parametrize("learner_cls", PREDICT_REORDERED) +def test_predict_warns_positionally(generate_regression_data, learner_cls): + """After the flip fit and predict must agree, so predict warns too.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = learner_cls(learner=LinearRegression()) + learner.fit(X=X, y=y, treatment=treatment) + with pytest.warns(FutureWarning, match="argument order"): + learner.predict(X, treatment, y) + + +@pytest.mark.parametrize("learner_cls", PREDICT_REORDERED) +def test_predict_keyword_is_silent(generate_regression_data, learner_cls): + y, X, treatment, _, _, _ = generate_regression_data() + learner = learner_cls(learner=LinearRegression()) + learner.fit(X=X, y=y, treatment=treatment) + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.predict(X, treatment=treatment, y=y) + assert _order_warnings(record) == [] + + +def test_fit_predict_warns_once_across_fit_and_predict(generate_regression_data): + """fit_predict delegates to both fit and predict; still one warning.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = BaseTRegressor(learner=LinearRegression()) + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + learner.fit_predict(X, treatment, y) + assert len(_order_warnings(record)) == 1 + + +# --- classifier fit overrides are wrapped independently --------------------- + + +@pytest.mark.parametrize( + "learner_cls", [BaseSClassifier, BaseTClassifier, BaseXClassifier, BaseRClassifier] +) +def test_classifier_fit_is_shimmed(learner_cls): + """The classifier variants define their own fit, so each is wrapped separately.""" + assert getattr(learner_cls.fit, "_arg_order_shimmed", False) + assert _positional_params(learner_cls.fit)[:3] == ["X", "treatment", "y"] + + +# --- #981: coverage beyond the fit family ----------------------------------- + + +def test_bootstrap_warns_positionally(generate_regression_data): + """`bootstrap` is public API on the most-used learners and changes at v1.0.""" + y, X, treatment, _, _, _ = generate_regression_data() + learner = BaseTRegressor(learner=LinearRegression()) + learner.fit(X=X, y=y, treatment=treatment) + with pytest.warns(FutureWarning, match="argument order"): + learner.bootstrap(X, treatment, y, None, 200) + + +def test_sensitivity_get_prediction_message_moves_p_right(): + """Sensitivity takes (X, p, treatment, y) — a third shape (#980 q2).""" + params = _positional_params(Sensitivity.get_prediction.__wrapped__) + assert params == ["X", "p", "treatment", "y"] + assert v1_order(params) == ["X", "y", "treatment", "p"] + + +def test_uplift_tree_validation_triple_is_reordered(): + """#980 q1: X_val/y_val/treatment_val stays consistent with the main triple.""" + params = _positional_params(UpliftTreeClassifier.fit.__wrapped__) + assert v1_order(params) == [ + "X", + "y", + "treatment", + "X_val", + "y_val", + "treatment_val", + "sample_weight", + "check_input", + ] + + +def test_timeit_preserves_signature(): + """`bootstrap_pool` is behind @timeit; without functools.wraps its signature + reads (*args, **kw), which hid it from the signature-based shim.""" + assert CausalTreeRegressor.bootstrap_pool.__name__ == "bootstrap_pool" + params = _positional_params(CausalTreeRegressor.bootstrap_pool.__wrapped__) + assert params[:3] == ["X", "treatment", "y"]