Add --color CLI flag, supporting "auto", "never", and "always" (#227)

* Use `--color=WHEN` style for color CLI argument

* Fix failing pre-commit check

* Address review suggestions

- Use `click.Context.color` as the single source of truth for the color setting.
- Set the color setting using a callback.
- Add type hints to `tests/unit/test_cli_parse.py`.
- Update text color tests to capture the Click context, and verify that its `color` attribute has the expected value.
- Replaced `TextReporter._style` with `click.style`

Author: WillDaSilva

src/check_jsonschema/cli.py

@@ -28,11 +28,6 @@
 )
 
 
-def evaluate_environment_settings(ctx: click.Context) -> None:
-    if os.getenv("NO_COLOR") is not None:
-        ctx.color = False
-
-
 class SchemaLoadingMode(enum.Enum):
     filepath = "filepath"
     builtin = "builtin"
@@ -95,6 +90,17 @@ def format_opts(self) -> FormatOptions:
         )
 
 
+def set_color_mode(ctx: click.Context, param: str, value: str) -> None:
+    if "NO_COLOR" in os.environ:
+        ctx.color = False
+    else:
+        ctx.color = {
+            "auto": None,
+            "always": True,
+            "never": False,
+        }[value]
+
+
 @click.command(
     "check-jsonschema",
     help="""\
@@ -211,6 +217,14 @@ def format_opts(self) -> FormatOptions:
     type=click.Choice(tuple(REPORTER_BY_NAME.keys()), case_sensitive=False),
     default="text",
 )
+@click.option(
+    "--color",
+    help="Force or disable colorized output. Defaults to autodetection.",
+    default="auto",
+    type=click.Choice(("auto", "always", "never")),
+    callback=set_color_mode,
+    expose_value=False,
+)
 @click.option(
     "-v",
     "--verbose",
@@ -227,9 +241,7 @@ def format_opts(self) -> FormatOptions:
     count=True,
 )
 @click.argument("instancefiles", required=True, nargs=-1)
-@click.pass_context
 def main(
-    ctx: click.Context,
     *,
     schemafile: str | None,
     builtin_schema: str | None,
@@ -248,7 +260,6 @@ def main(
     instancefiles: tuple[str, ...],
 ) -> None:
     args = ParseResult()
-    evaluate_environment_settings(ctx)
 
     args.set_schema(schemafile, builtin_schema, check_metaschema)
     args.instancefiles = instancefiles

src/check_jsonschema/reporter.py

@@ -44,25 +44,17 @@ def __init__(
         *,
         verbosity: int,
         stream: t.TextIO | None = None,  # default stream is stdout (None)
-        color: bool = True,
     ) -> None:
         super().__init__(verbosity=verbosity)
         self.stream = stream
-        self.color = color
 
     def _echo(self, s: str, *, indent: int = 0) -> None:
         click.echo(" " * indent + s, file=self.stream)
 
-    def _style(self, s: str, *, fg: str | None = None) -> str:
-        if self.color:
-            return click.style(s, fg=fg)
-        else:
-            return s
-
     def report_success(self) -> None:
         if self.verbosity < 1:
             return
-        ok = self._style("ok", fg="green")
+        ok = click.style("ok", fg="green")
         self._echo(f"{ok} -- validation done")
 
     def _format_validation_error_message(
@@ -71,8 +63,7 @@ def _format_validation_error_message(
         error_loc = err.json_path
         if filename:
             error_loc = f"{filename}::{error_loc}"
-        if self.color:
-            error_loc = self._style(error_loc, fg="yellow")
+        error_loc = click.style(error_loc, fg="yellow")
         return f"{error_loc}: {err.message}"
 
     def _show_validation_error(
@@ -95,7 +86,7 @@ def _show_validation_error(
 
     def _show_parse_error(self, filename: str, err: ParseError) -> None:
         if self.verbosity < 2:
-            self._echo(self._style(str(err), fg="yellow"), indent=2)
+            self._echo(click.style(str(err), fg="yellow"), indent=2)
         elif self.verbosity < 3:
             self._echo(textwrap.indent(format_error(err, mode="short"), "  "))
         else:

tests/unit/test_cli_parse.py

@@ -1,14 +1,15 @@
+from __future__ import annotations
+
+import typing as t
 from unittest import mock
 
+import click
 import pytest
 from click.testing import CliRunner
+from click.testing import Result as ClickResult
 
 from check_jsonschema import main as cli_main
-from check_jsonschema.cli import (
-    ParseResult,
-    SchemaLoadingMode,
-    evaluate_environment_settings,
-)
+from check_jsonschema.cli import ParseResult, SchemaLoadingMode
 
 
 @pytest.fixture
@@ -26,10 +27,32 @@ def mock_cli_exec():
 
 
 @pytest.fixture
-def runner():
+def runner() -> CliRunner:
     return CliRunner(mix_stderr=False)
 
 
+def invoke_and_get_ctx(
+    runner: CliRunner,
+    cmd: click.Command,
+    args: t.Sequence[str],
+) -> tuple[ClickResult, click.Context]:
+    # There doesn't appear to be a good way to get the Click context used by a
+    # test invocation, so we replace the invoke method with a wrapper that
+    # calls `click.get_current_context` to extract the context object.
+
+    ctx = None
+
+    def extract_ctx(*args, **kwargs):
+        nonlocal ctx
+        ctx = click.get_current_context()
+        return click.Command.invoke(*args, **kwargs)
+
+    with mock.patch("click.Command.invoke", extract_ctx):
+        results = runner.invoke(cmd, args)
+
+    return results, ctx
+
+
 @pytest.mark.parametrize(
     "schemafile,builtin_schema,check_metaschema,expect_mode",
     [
@@ -55,34 +78,34 @@ def test_parse_result_set_schema(
         assert args.schema_path is None
 
 
-def test_requires_some_args(runner):
+def test_requires_some_args(runner: CliRunner):
     result = runner.invoke(cli_main, [])
     assert result.exit_code == 2
 
 
-def test_schemafile_and_instancefile(runner, mock_parse_result):
+def test_schemafile_and_instancefile(runner: CliRunner, mock_parse_result: ParseResult):
     runner.invoke(cli_main, ["--schemafile", "schema.json", "foo.json"])
     assert mock_parse_result.schema_mode == SchemaLoadingMode.filepath
     assert mock_parse_result.schema_path == "schema.json"
     assert mock_parse_result.instancefiles == ("foo.json",)
 
 
-def test_requires_at_least_one_instancefile(runner):
+def test_requires_at_least_one_instancefile(runner: CliRunner):
     result = runner.invoke(cli_main, ["--schemafile", "schema.json"])
     assert result.exit_code == 2
 
 
-def test_requires_schemafile(runner):
+def test_requires_schemafile(runner: CliRunner):
     result = runner.invoke(cli_main, ["foo.json"])
     assert result.exit_code == 2
 
 
-def test_no_cache_defaults_false(runner, mock_parse_result):
+def test_no_cache_defaults_false(runner: CliRunner, mock_parse_result: ParseResult):
     runner.invoke(cli_main, ["--schemafile", "schema.json", "foo.json"])
     assert mock_parse_result.disable_cache is False
 
 
-def test_no_cache_flag_is_true(runner, mock_parse_result):
+def test_no_cache_flag_is_true(runner: CliRunner, mock_parse_result: ParseResult):
     runner.invoke(cli_main, ["--schemafile", "schema.json", "foo.json", "--no-cache"])
     assert mock_parse_result.disable_cache is True
 
@@ -119,7 +142,7 @@ def test_no_cache_flag_is_true(runner, mock_parse_result):
         ],
     ],
 )
-def test_mutex_schema_opts(runner, cmd_args):
+def test_mutex_schema_opts(runner: CliRunner, cmd_args: list[str]):
     result = runner.invoke(cli_main, cmd_args)
     assert result.exit_code == 2
     assert "are mutually exclusive" in result.stderr
@@ -133,22 +156,68 @@ def test_mutex_schema_opts(runner, cmd_args):
         ["-h"],
     ],
 )
-def test_supports_common_option(runner, cmd_args):
+def test_supports_common_option(runner: CliRunner, cmd_args: list[str]):
     result = runner.invoke(cli_main, cmd_args)
     assert result.exit_code == 0
 
 
 @pytest.mark.parametrize(
     "setting,expect_value", [(None, None), ("1", False), ("0", False)]
 )
-def test_no_color_env_var(monkeypatch, setting, expect_value):
-    mock_ctx = mock.Mock()
-    mock_ctx.color = None
-
+def test_no_color_env_var(
+    runner: CliRunner,
+    monkeypatch: pytest.MonkeyPatch,
+    setting: str | None,
+    expect_value: bool | None,
+):
     if setting is None:
         monkeypatch.delenv("NO_COLOR", raising=False)
     else:
         monkeypatch.setenv("NO_COLOR", setting)
 
-    evaluate_environment_settings(mock_ctx)
-    assert mock_ctx.color == expect_value
+    _, ctx = invoke_and_get_ctx(
+        runner, cli_main, ["--schemafile", "schema.json", "foo.json"]
+    )
+    assert ctx.color == expect_value
+
+
+@pytest.mark.parametrize(
+    "setting,expected_value",
+    [(None, None), ("auto", None), ("always", True), ("never", False)],
+)
+def test_color_cli_option(
+    runner: CliRunner,
+    setting: str | None,
+    expected_value: bool | None,
+):
+    args = ["--schemafile", "schema.json", "foo.json"]
+    if setting:
+        args.extend(("--color", setting))
+    _, ctx = invoke_and_get_ctx(runner, cli_main, args)
+    assert ctx.color == expected_value
+
+
+def test_no_color_env_var_overrides_cli_option(
+    runner: CliRunner, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("NO_COLOR", "1")
+
+    _, ctx = invoke_and_get_ctx(
+        runner, cli_main, ["--color=always", "--schemafile", "schema.json", "foo.json"]
+    )
+    assert ctx.color is False
+
+
+@pytest.mark.parametrize(
+    "setting,expected_value",
+    [("auto", 0), ("always", 0), ("never", 0), ("anything_else", 2)],
+)
+def test_color_cli_option_is_choice(
+    runner: CliRunner, setting: str, expected_value: int
+):
+    assert (
+        runner.invoke(
+            cli_main, ["--color", setting, "--schemafile", "schema.json", "foo.json"]
+        ).exit_code
+        == expected_value
+    )

tests/unit/test_reporters.py

@@ -14,7 +14,7 @@ def _make_success_result():
 @pytest.mark.parametrize("verbosity", (0, 1, 2))
 @pytest.mark.parametrize("use_report_result_path", (False, True))
 def test_text_format_success(capsys, verbosity, use_report_result_path):
-    reporter = TextReporter(color=False, verbosity=verbosity)
+    reporter = TextReporter(verbosity=verbosity)
     if use_report_result_path:
         reporter.report_result(_make_success_result())
     else:
@@ -58,14 +58,18 @@ def test_text_format_validation_error_message_simple():
     )
     err = next(validator.iter_errors({"foo": {"bar": 1}}))
 
-    text_reporter = TextReporter(color=False, verbosity=1)
+    text_reporter = TextReporter(verbosity=1)
     s1 = text_reporter._format_validation_error_message(err, filename="foo.json")
-    assert (
-        s1 == "foo.json::$.foo: {'bar': 1} is not valid under any of the given schemas"
+    assert s1 == (
+        "\x1b[33mfoo.json::$.foo\x1b[0m: {'bar': 1} "
+        "is not valid under any of the given schemas"
     )
 
     s2 = text_reporter._format_validation_error_message(err)
-    assert s2 == "$.foo: {'bar': 1} is not valid under any of the given schemas"
+    assert s2 == (
+        "\x1b[33m$.foo\x1b[0m: {'bar': 1} "
+        "is not valid under any of the given schemas"
+    )
 
 
 @pytest.mark.parametrize("verbosity", (0, 1, 2))
@@ -104,7 +108,7 @@ def test_text_print_validation_error_nested(capsys, verbosity):
     result = CheckResult()
     result.record_validation_error("foo.json", err)
 
-    text_reporter = TextReporter(color=False, verbosity=verbosity)
+    text_reporter = TextReporter(verbosity=verbosity)
     text_reporter.report_result(result)
     captured = capsys.readouterr()
     # nothing to stderr