Log env packages

fancylog/fancylog.py

@@ -1,8 +1,10 @@
 """Wrapper around the standard logging module, with additional information."""
 
 import contextlib
+import json
 import logging
 import os
+import subprocess
 import sys
 from datetime import datetime
 from importlib.util import find_spec
@@ -28,6 +30,8 @@ def start_logging(
     write_header=True,
     write_git=True,
     write_cli_args=True,
+    write_python=True,
+    write_env_packages=True,
     write_variables=True,
     log_to_file=True,
     log_to_console=True,
@@ -62,6 +66,10 @@ def start_logging(
         Write information about the git repository. Default: True
     write_cli_args
         Log the command-line arguments. Default: True
+    write_python
+        Log the Python version. Default: True
+    write_env_packages
+        Log the packages in the environment. Default: True
     write_variables
         Write the attributes of selected objects. Default: True
     log_to_file
@@ -105,6 +113,8 @@ def start_logging(
             write_header=write_header,
             write_git=write_git,
             write_cli_args=write_cli_args,
+            write_python=write_python,
+            write_env_packages=write_env_packages,
             write_variables=write_variables,
             log_header=log_header,
         )
@@ -132,6 +142,8 @@ def __init__(
         write_header=True,
         write_git=True,
         write_cli_args=True,
+        write_python=True,
+        write_env_packages=True,
         write_variables=True,
         log_header=None,
     ):
@@ -148,6 +160,10 @@ def __init__(
                 self.write_git_info(self.program.__name__)
             if write_cli_args:
                 self.write_command_line_arguments()
+            if write_python:
+                self.write_python_version()
+            if write_env_packages:
+                self.write_environment_packages()
             if write_variables and variable_objects:
                 self.write_variables(variable_objects)
 
@@ -206,6 +222,105 @@ def write_command_line_arguments(self, header="COMMAND LINE ARGUMENTS"):
         self.file.write(f"Command: {sys.argv[0]} \n")
         self.file.write(f"Input arguments: {sys.argv[1:]}")
 
+    def write_python_version(self, header="PYTHON VERSION"):
+        """Write the Python version used to run the script.
+
+        Parameters
+        ----------
+        header
+            Title of the section that will be written to the log file.
+
+        """
+        self.write_separated_section_header(header)
+        self.file.write(f"Python version: {sys.version.split()[0]}")
+
+    def write_environment_packages(self, header="ENVIRONMENT"):
+        """Write the local/global environment packages used to run the script.
+
+        Attempt to collect conda packages and, if this fails,
+        collect pip packages.
+
+        Parameters
+        ----------
+        header
+            Title of the section that will be written to the log file
+
+        """
+        self.write_separated_section_header(header)
+
+        # Attempt to log conda env name and packages
+        try:
+            conda_env = os.environ["CONDA_PREFIX"].split(os.sep)[-1]
+            conda_exe = os.environ["CONDA_EXE"]
+            conda_list = subprocess.run(
+                [conda_exe, "list", "--json"], capture_output=True, text=True
+            )
+
+            env_pkgs = json.loads(conda_list.stdout)
+
+            self.file.write(f"Conda environment: {conda_env}\n\n")
+            self.file.write("Environment packages (conda):\n")
+            self.write_packages(env_pkgs)
+
+        # If no conda env, fall back to logging pip
+        except KeyError:
+            python_executable = sys.executable
+            pip_list = subprocess.run(
+                [
+                    python_executable,
+                    "-m",
+                    "pip",
+                    "list",
+                    "--verbose",
+                    "--format=json",
+                ],
+                capture_output=True,
+                text=True,
+            )
+
+            all_pkgs = json.loads(pip_list.stdout)
+
+            try:
+                # If there is a local env, log local packages first
+                env_pkgs = [
+                    pkg
+                    for pkg in all_pkgs
+                    if os.getenv("VIRTUAL_ENV") in pkg["location"]
+                ]
+
+                self.file.write(
+                    "No conda environment found, reporting pip packages\n\n"
+                )
+                self.file.write("Local environment packages (pip):\n")
+                self.write_packages(env_pkgs)
+                self.file.write("\n")
+
+                # Log global-available packages (if any)
+                global_pkgs = [pkg for pkg in all_pkgs if pkg not in env_pkgs]
+
+                self.file.write("Global environment packages (pip):\n")
+                self.write_packages(global_pkgs)
+
+            except TypeError:
+                self.file.write(
+                    "No environment found, reporting global pip packages\n\n"
+                )
+                self.write_packages(all_pkgs)
+
+    def write_packages(self, env_pkgs):
+        """Write the packages in the local environment.
+
+        Parameters
+        ----------
+        env_pkgs
+            A dictionary of environment packages, the name and version
+            of which will be written.
+
+        """
+        self.file.write(f"{'Name':20} {'Version':15}\n")
+        for pkg in env_pkgs:
+            self.file.write(f"{pkg['name']:20} {pkg['version']:15}\n")
+
     def write_variables(self, variable_objects):
         """Write a section for variables with their values.
 

tests/tests/test_general.py

@@ -1,4 +1,11 @@
+import json
 import logging
+import os
+import platform
+import subprocess
+import sys
+from importlib.metadata import distributions
+from unittest.mock import MagicMock, patch
 
 import pytest
 from rich.logging import RichHandler
@@ -206,3 +213,215 @@ def test_named_logger_doesnt_propagate(tmp_path, capsys):
     assert "PQ&*" not in captured.out, (
         "logger writing to stdout through root handler"
     )
+
+
+@pytest.mark.parametrize("boolean, operator", [(True, True), (False, False)])
+def test_python_version_header(boolean, operator, tmp_path):
+    ver_header = f"{lateral_separator}  PYTHON VERSION  {lateral_separator}\n"
+
+    fancylog.start_logging(tmp_path, fancylog, write_python=boolean)
+
+    log_file = next(tmp_path.glob("*.log"))
+
+    with open(log_file) as file:
+        assert (ver_header in file.read()) == operator
+
+
+def test_correct_python_version_logged(tmp_path):
+    """Python version logged should be equal to
+    the output of platform.python_version().
+    """
+
+    fancylog.start_logging(tmp_path, fancylog, write_python=True)
+
+    log_file = next(tmp_path.glob("*.log"))
+
+    # Test logged python version is equal to platform.python_version()
+    with open(log_file) as file:
+        assert f"Python version: {platform.python_version()}" in file.read()
+
+
+@pytest.mark.parametrize("boolean, operator", [(True, True), (False, False)])
+def test_environment_header(boolean, operator, tmp_path):
+    ver_header = f"{lateral_separator}  ENVIRONMENT  {lateral_separator}\n"
+
+    fancylog.start_logging(tmp_path, fancylog, write_env_packages=boolean)
+
+    log_file = next(tmp_path.glob("*.log"))
+
+    with open(log_file) as file:
+        assert (ver_header in file.read()) == operator
+
+
+def test_correct_pkg_version_logged(tmp_path):
+    """Package versions logged should be equal to
+    the output of `conda list` or `pip list`.
+    """
+    fancylog.start_logging(tmp_path, fancylog, write_env_packages=True)
+
+    log_file = next(tmp_path.glob("*.log"))
+
+    try:
+        # If there is a conda environment, assert that the correct
+        # version is logged for all pkgs
+        conda_exe = os.environ["CONDA_EXE"]
+        conda_list = subprocess.run(
+            [conda_exe, "list", "--json"], capture_output=True, text=True
+        )
+
+        conda_pkgs = json.loads(conda_list.stdout)
+        for pkg in conda_pkgs:
+            assert f"{pkg['name']:20} {pkg['version']:15}\n"
+
+    except KeyError:
+        # If there is no conda environment, assert that the correct
+        # version is logged for all packages logged with pip list
+        with open(log_file) as file:
+            file_content = file.read()
+
+            # Test local environment versions
+            local_site_packages = next(
+                p for p in sys.path if "site-packages" in p
+            )
+
+            for dist in distributions():
+                if str(dist.locate_file("")).startswith(local_site_packages):
+                    assert (
+                        f"{dist.metadata['Name']:20} {dist.version}"
+                        in file_content
+                    )
+
+
+def test_mock_pip_pkgs(tmp_path):
+    """Mock pip list subprocess
+    and test that packages are logged correctly.
+    """
+
+    # Simulated `pip list --json` output
+    fake_pip_output = json.dumps(
+        [
+            {"name": "fancylog", "version": "1.1.1", "location": "fake_env"},
+            {"name": "pytest", "version": "1.1.1", "location": "global_env"},
+        ]
+    )
+
+    # Patch the environment and subprocess
+    with (
+        patch.dict(os.environ, {}, clear=False),
+        patch("os.getenv") as mock_getenv,
+        patch("subprocess.run") as mock_run,
+    ):
+        # Eliminate conda environment packages triggers logging pip list
+        os.environ.pop("CONDA_PREFIX", None)
+        os.environ.pop("CONDA_EXE", None)
+
+        mock_getenv.return_value = "fake_env"
+
+        # Mocked subprocess result
+        mock_run.return_value = MagicMock(stdout=fake_pip_output, returncode=0)
+
+        fancylog.start_logging(tmp_path, fancylog, write_env_packages=True)
+
+        log_file = next(tmp_path.glob("*.log"))
+
+        # Log contains conda subheaders and mocked pkgs versions
+        with open(log_file) as file:
+            file_content = file.read()
+
+            assert (
+                "No conda environment found, reporting pip packages"
+            ) in file_content
+
+            assert f"{'fancylog':20} {'1.1.1'}"
+            assert f"{'pytest':20} {'1.1.1'}"
+
+
+def test_mock_conda_pkgs(tmp_path):
+    """Mock conda environment variables
+    and test that packages are logged correctly.
+    """
+
+    fake_conda_env_name = "test_env"
+    fake_conda_prefix = os.path.join(
+        "path", "conda", "envs", fake_conda_env_name
+    )
+    fake_conda_exe = os.path.join("fake", "conda")
+
+    # Simulated `conda list --json` output
+    fake_conda_output = json.dumps(
+        [
+            {"name": "fancylog", "version": "1.1.1"},
+            {"name": "pytest", "version": "1.1.1"},
+        ]
+    )
+
+    # Patch the environment and subprocess
+    with (
+        patch.dict(
+            os.environ,
+            {"CONDA_PREFIX": fake_conda_prefix, "CONDA_EXE": fake_conda_exe},
+        ),
+        patch("subprocess.run") as mock_run,
+    ):
+        # Mocked subprocess result
+        mock_run.return_value = MagicMock(
+            stdout=fake_conda_output, returncode=0
+        )
+
+        fancylog.start_logging(tmp_path, fancylog, write_env_packages=True)
+
+        log_file = next(tmp_path.glob("*.log"))
+
+        # Log contains conda subheaders and mocked pkgs versions
+        with open(log_file) as file:
+            file_content = file.read()
+
+            assert "Conda environment:" in file_content
+            assert "Environment packages (conda):" in file_content
+            assert f"{'fancylog':20} {'1.1.1'}"
+            assert f"{'pytest':20} {'1.1.1'}"
+
+
+def test_mock_no_environment(tmp_path):
+    """Mock lack of any environment,
+    and test that packages are logged correctly.
+    """
+
+    # Simulated `pip list --json` output
+    fake_pip_output = json.dumps(
+        [
+            {"name": "fancylog", "version": "1.1.1", "location": "fake_env"},
+            {"name": "pytest", "version": "1.1.1", "location": "global_env"},
+        ]
+    )
+
+    # Patch the environment and subprocess
+    with (
+        patch.dict(os.environ, {}, clear=False),
+        patch("os.getenv") as mock_getenv,
+        patch("subprocess.run") as mock_run,
+    ):
+        # Eliminate conda environment packages triggers logging pip list
+        os.environ.pop("CONDA_PREFIX", None)
+        os.environ.pop("CONDA_EXE", None)
+
+        # Mock lack of any local environment
+        mock_getenv.return_value = None
+
+        # Mocked subprocess result
+        mock_run.return_value = MagicMock(stdout=fake_pip_output, returncode=0)
+
+        fancylog.start_logging(tmp_path, fancylog, write_env_packages=True)
+
+        log_file = next(tmp_path.glob("*.log"))
+
+        # Log contains conda subheaders and mocked pkgs versions
+        with open(log_file) as file:
+            file_content = file.read()
+
+            assert (
+                "No environment found, reporting global pip packages"
+            ) in file_content
+
+            assert f"{'fancylog':20} {'1.1.1'}"
+            assert f"{'pytest':20} {'1.1.1'}"