From f3f1d9069baf2704507f2684f2131d5e469be70a Mon Sep 17 00:00:00 2001 From: Keith Battocchi Date: Fri, 20 Oct 2023 22:14:59 -0400 Subject: [PATCH] Remove legacy assertWarns hack Signed-off-by: Keith Battocchi --- econml/tests/test_automated_ml.py | 1 - econml/tests/test_dml.py | 1 - econml/tests/test_drlearner.py | 1 - econml/tests/test_dynamic_dml.py | 1 - econml/tests/test_metalearners.py | 1 - econml/tests/test_random_state.py | 1 - econml/tests/utilities.py | 22 ---------------------- 7 files changed, 28 deletions(-) diff --git a/econml/tests/test_automated_ml.py b/econml/tests/test_automated_ml.py index ba222fad9..1e30349b0 100644 --- a/econml/tests/test_automated_ml.py +++ b/econml/tests/test_automated_ml.py @@ -22,7 +22,6 @@ import itertools from econml.sklearn_extensions.linear_model import WeightedLasso from econml.tests.test_statsmodels import _summarize -import econml.tests.utilities # bugfix for assertWarns import copy import logging from econml.data.dgps import ihdp_surface_B diff --git a/econml/tests/test_dml.py b/econml/tests/test_dml.py index d33c8b418..2ea12b851 100644 --- a/econml/tests/test_dml.py +++ b/econml/tests/test_dml.py @@ -19,7 +19,6 @@ import itertools from econml.sklearn_extensions.linear_model import WeightedLasso, StatsModelsRLM, StatsModelsLinearRegression from econml.tests.test_statsmodels import _summarize -import econml.tests.utilities # bugfix for assertWarns from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier from sklearn.multioutput import MultiOutputRegressor from econml.grf import MultiOutputGRF diff --git a/econml/tests/test_drlearner.py b/econml/tests/test_drlearner.py index c1b29cb77..7b84f6a62 100644 --- a/econml/tests/test_drlearner.py +++ b/econml/tests/test_drlearner.py @@ -22,7 +22,6 @@ from econml.inference import BootstrapInference, StatsModelsInferenceDiscrete from econml.utilities import get_feature_names_or_default, shape, hstack, vstack, reshape, cross_product from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression -import econml.tests.utilities # bugfix for assertWarns from econml.tests.utilities import (GroupingModel, NestedModel) diff --git a/econml/tests/test_dynamic_dml.py b/econml/tests/test_dynamic_dml.py index 0c2c7339f..118f51563 100644 --- a/econml/tests/test_dynamic_dml.py +++ b/econml/tests/test_dynamic_dml.py @@ -12,7 +12,6 @@ from econml.panel.dml._dml import _get_groups_period_filter from econml.inference import BootstrapInference, EmpiricalInferenceResults, NormalInferenceResults from econml.utilities import shape, hstack, vstack, reshape, cross_product -import econml.tests.utilities # bugfix for assertWarns from econml.tests.dgp import DynamicPanelDGP diff --git a/econml/tests/test_metalearners.py b/econml/tests/test_metalearners.py index 534f1623b..7c6050641 100644 --- a/econml/tests/test_metalearners.py +++ b/econml/tests/test_metalearners.py @@ -9,7 +9,6 @@ from sklearn.pipeline import Pipeline from sklearn.preprocessing import PolynomialFeatures from econml.metalearners import * -import econml.tests.utilities # bugfix for assertWarns class TestMetalearners(unittest.TestCase): diff --git a/econml/tests/test_random_state.py b/econml/tests/test_random_state.py index cd8f27265..8fe45ebce 100644 --- a/econml/tests/test_random_state.py +++ b/econml/tests/test_random_state.py @@ -18,7 +18,6 @@ import itertools from econml.sklearn_extensions.linear_model import WeightedLasso, StatsModelsRLM, StatsModelsLinearRegression from econml.tests.test_statsmodels import _summarize -import econml.tests.utilities # bugfix for assertWarns from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier from econml.iv.dr._dr import _DummyCATE diff --git a/econml/tests/utilities.py b/econml/tests/utilities.py index db9ac1f59..4c04cc89d 100644 --- a/econml/tests/utilities.py +++ b/econml/tests/utilities.py @@ -1,28 +1,6 @@ # Copyright (c) PyWhy contributors. All rights reserved. # Licensed under the MIT License. -import unittest -import sys -import warnings - -# HACK: work around bug in assertWarns (https://bugs.python.org/issue29620) -# this can be removed if the corresponding pull request (https://github.com/python/cpython/pull/4800) is ever merged - - -def _enter(self): - # The __warningregistry__'s need to be in a pristine state for tests - # to work properly. - for v in list(sys.modules.values()): - if getattr(v, '__warningregistry__', None): - v.__warningregistry__ = {} - self.warnings_manager = warnings.catch_warnings(record=True) - self.warnings = self.warnings_manager.__enter__() - warnings.simplefilter("always", self.expected) - return self - - -unittest.case._AssertWarnsContext.__enter__ = _enter - import numpy as np from sklearn.calibration import check_cv