diff --git a/experiment_utils/experiment_analyzer.py b/experiment_utils/experiment_analyzer.py index 731a834..4fceb69 100644 --- a/experiment_utils/experiment_analyzer.py +++ b/experiment_utils/experiment_analyzer.py @@ -14,7 +14,7 @@ from scipy.stats import gaussian_kde from pyspark.sql import functions as F from pyspark.sql import DataFrame -from .utils import turn_off_package_logger, log_and_raise_error +from .utils import turn_off_package_logger, log_and_raise_error, get_logger from .spark_instance import * @@ -66,18 +66,7 @@ def __init__( List of covariates to include in the final linear regression model, by default None """ - self.logger = logging.getLogger('Experiment Analyzer') - self.logger.setLevel(logging.INFO) - if self.logger.hasHandlers(): - self.logger.handlers.clear() - console_handler = logging.StreamHandler() - console_handler.setFormatter( - logging.Formatter( - fmt='%(asctime)s - %(message)s', - datefmt='%d/%m/%Y %I:%M:%S %p' - ) - ) - self.logger.addHandler(console_handler) + self.logger = get_logger('Experiment Analyzer') self.data = self.__ensure_spark_df(data) self.outcomes = self.__ensure_list(outcomes) self.covariates = self.__ensure_list(covariates) diff --git a/experiment_utils/utils.py b/experiment_utils/utils.py index 38a59b4..b07b6c1 100644 --- a/experiment_utils/utils.py +++ b/experiment_utils/utils.py @@ -17,6 +17,26 @@ def turn_off_package_logger(package: str): logger.handlers = [logging.NullHandler()] +def get_logger(name: str): + logger = logging.getLogger(name) + + if not logger.hasHandlers(): + logger.setLevel(logging.INFO) + + while logger.hasHandlers(): + logger.removeHandler(logger.handlers[0]) + + console_handler = logging.StreamHandler() + formatter = logging.Formatter( + fmt='%(asctime)s - %(message)s', + datefmt='%d/%m/%Y %I:%M:%S %p' + ) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + return logger + + def log_and_raise_error(logger, message, exception_type=ValueError): """" Logs an error message and raises an exception of the specified type.