diff --git a/src/muse/__init__.py b/src/muse/__init__.py index 3f6b24984..a0fe9856a 100644 --- a/src/muse/__init__.py +++ b/src/muse/__init__.py @@ -34,6 +34,49 @@ def _create_logger(color: bool = True): return logger +def add_file_logger() -> None: + """Adds a file logger to the main logger. + + The file logger is split into two files: one for INFO and DEBUG messages, and one + for WARNING messages and above to avoid cluttering the main log file and highlight + potential issues. + """ + import logging + from pathlib import Path + + from .defaults import DEFAULT_OUTPUT_DIRECTORY + + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + + DEFAULT_OUTPUT_DIRECTORY.mkdir(parents=True, exist_ok=True) + + # Sets the warning log, for warnings and above + warning_file = Path(DEFAULT_OUTPUT_DIRECTORY) / "muse_warning.log" + if warning_file.exists(): + warning_file.unlink() + + warning_file_handler = logging.FileHandler(warning_file) + warning_file_handler.setLevel(logging.WARNING) + warning_file_handler.setFormatter(formatter) + warning_file_handler.filters = [lambda record: record.levelno > logging.INFO] + + logging.getLogger("muse").addHandler(warning_file_handler) + + # Sets the info log, for debug and info only + info_file = Path(DEFAULT_OUTPUT_DIRECTORY) / "muse_info.log" + if info_file.exists(): + info_file.unlink() + + info_file_handler = logging.FileHandler(info_file) + info_file_handler.setLevel(logging.DEBUG) + info_file_handler.setFormatter(formatter) + info_file_handler.filters = [lambda record: record.levelno <= logging.INFO] + + logging.getLogger("muse").addHandler(info_file_handler) + + logger = _create_logger(os.environ.get("MUSE_COLOR_LOG") != "False") """ Main logger """ diff --git a/src/muse/__main__.py b/src/muse/__main__.py index a71bc7f0a..6108b0b8a 100644 --- a/src/muse/__main__.py +++ b/src/muse/__main__.py @@ -38,7 +38,7 @@ def muse_main(settings, model, copy): from logging import getLogger from pathlib import Path - from muse import examples + from muse import add_file_logger, examples from muse.mca import MCA from muse.readers.toml import read_settings @@ -53,6 +53,7 @@ def muse_main(settings, model, copy): else: settings = read_settings(settings) getLogger("muse").setLevel(settings.log_level) + add_file_logger() MCA.factory(settings).run() diff --git a/src/muse/examples.py b/src/muse/examples.py index 2ccc37d59..9aa4fa9b6 100644 --- a/src/muse/examples.py +++ b/src/muse/examples.py @@ -51,10 +51,21 @@ def available_examples() -> list[str]: return [d.stem for d in example_data_dir().iterdir() if d.is_dir()] -def model(name: str = "default") -> MCA: - """Fully constructs a given example model.""" +def model(name: str = "default", test: bool = False) -> MCA: + """Fully constructs a given example model. + + File logging is added if ``test`` is False. + + Args: + name: Name of the model to load. + test: If True, the logging to file is not added. + + Returns: + The MCA model. + """ from tempfile import TemporaryDirectory + from muse import add_file_logger from muse.readers.toml import read_settings # we could modify the settings directly, but instead we use the copy_model function. @@ -63,6 +74,8 @@ def model(name: str = "default") -> MCA: path = copy_model(name, tmpdir) settings = read_settings(path / "settings.toml") getLogger("muse").setLevel(settings.log_level) + if not test: + add_file_logger() return MCA.factory(settings) diff --git a/tests/test_aggregoutput.py b/tests/test_aggregoutput.py index a5578d19f..f11e8e74b 100644 --- a/tests/test_aggregoutput.py +++ b/tests/test_aggregoutput.py @@ -9,7 +9,7 @@ def test_aggregate_sector(): """ from pandas import DataFrame, concat - mca = examples.model("multiple_agents") + mca = examples.model("multiple_agents", test=True) year = [2020, 2025] sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] agent_list = [list(a.agents) for a in sector_list] @@ -43,7 +43,7 @@ def test_aggregate_sectors(): from muse.outputs.mca import _aggregate_sectors - mca = examples.model("multiple_agents") + mca = examples.model("multiple_agents", test=True) year = [2020, 2025, 2030] sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] agent_list = [list(a.agents) for a in sector_list] @@ -83,7 +83,7 @@ def test_aggregate_sector_manyregions(): from muse.outputs.mca import _aggregate_sectors - mca = examples.model("multiple_agents") + mca = examples.model("multiple_agents", test=True) residential = next(sector for sector in mca.sectors if sector.name == "residential") agents = list(residential.agents) agents[0].assets["region"] = "BELARUS" diff --git a/tests/test_subsector.py b/tests/test_subsector.py index 9c326f1f4..244f80bd1 100644 --- a/tests/test_subsector.py +++ b/tests/test_subsector.py @@ -34,7 +34,7 @@ def test_subsector_investing_aggregation(): sector_list = ["residential", "power", "gas"] for model in model_list: - mca = examples.model(model) + mca = examples.model(model, test=True) for sname in sector_list: agents = list(examples.sector(sname, model).agents) sector = next(sector for sector in mca.sectors if sector.name == sname)