Repository: dtrifiro/pytest-fzf (language: python)

Issue:
add script to preview code without using bat
A script could be added to the package to print a syntax-highlighted version of the code using `pygments` and used as preview command for fzf


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index b6a35d4..57d039e 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,9 @@
 
 fzf-based test selection with `pytest`
 
-[![demo](https://github.com/dtrifiro/pytest-fzf/assets/36171005/29f7a610-2f15-402f-a24f-af8bf7e0e71d)](https://asciinema.org/a/iAr18ilruuPM7pZ1EAfXkxfEf)
+[![demo](https://github.com/dtrifiro/pytest-fzf/assets/36171005/d8a162fc-eed4-4382-9527-dc0cb58ed245)](https://asciinema.org/a/CfvBIUShAllMANUmXgYx0LYpM)
+
+(demo uses the [gruvbox-dark pygments style](https://pygments.org/styles/#gruvbox-dark))
 
 ---
 
@@ -17,7 +19,6 @@ fzf-based test selection with `pytest`
 ## Requirements
 
 - [fzf](https://github.com/junegunn/fzf)
-- (Optional, for colored preview of test functions) `bat`[sharkdp/bat](https://github.com/sharkdp/bat)
 
 ## Installation
 
@@ -32,6 +33,24 @@ pip install pytest-fzf
 
 ```bash
 pytest --fzf [query]
+pytest --fzf [--fzf-bat-preview] # uses bat as fzf preview command
+```
+
+### Syntax highlighting theme
+
+The theme used for previewing test functions can be set using `PYTEST_FZF_THEME` (or `BAT_THEME`, if you use `bat`, see [sharkdp/bat](https://github.com/sharkdp/bat)):
+
+```bash
+export PYTEST_FZF_THEME='gruvbox-dark'
+```
+
+For a list of supported themes, see https://pygments.org/styles/ or get a list by running:
+
+```python
+import pygments
+
+for style in pygments.styles.get_all_styles():
+    print(style)
 ```
 
 ### Keybindings
diff --git a/pyproject.toml b/pyproject.toml
index 4c6d63c..84f6788 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,8 +23,12 @@ dynamic = ["version"]
 dependencies = [
     "iterfzf==1.1.0.44.0",
     "pytest>=6.0.0",
+    "pygments>=2.15.0",
 ]
 
+[project.scripts]
+pytest-fzf-preview = "pytest_fzf.previewer:cli"
+
 [project.urls]
 Issues = "https://github.com/dtrifiro/pytest-fzf/issues"
 Source = "https://github.com/dtrifiro/pytest-fzf"
@@ -38,6 +42,7 @@ tests = [
     "pytest-cov==4.1.0",
     "pytest-mock==3.12.0",
     "mypy==1.8.0",
+    "types-pygments>=2.15.0",
 ]
 dev = [
     "pytest-fzf[tests]",
@@ -104,6 +109,13 @@ ignore = [
     "D100", # Missing docstring in public module
     "D103", # Missing docstring in public function
     "D104", # Missing docstring in public package
+
+    "EM101", # Exception must not use a string literal, assign to variable first
+
+    "D203", # one blank line before class
+    "D213", # multi-line-summary-second-line
+
+    "PTH123" # open() should be replaced by Path.open
 ]
 select = ["ALL"]
 show-source = true
@@ -112,6 +124,7 @@ show-fixes = true
 [tool.ruff.per-file-ignores]
 "noxfile.py" = ["D", "PTH"]
 "tests/**" = [ "S", "ARG001", "ARG002", "ANN"]
+"src/pytest_fzf/previewer.py" = ["T201"]
 
 [tool.ruff.lint.flake8-type-checking]
 strict = true


Test patch:
diff --git a/src/pytest_fzf/main.py b/src/pytest_fzf/main.py
index d628f26..c602578 100644
--- a/src/pytest_fzf/main.py
+++ b/src/pytest_fzf/main.py
@@ -8,7 +8,6 @@
 BAT_AVAILABLE = shutil.which("bat")
 BAT_CMD = "bat --color=always --language=python"
 
-
 _sentinel = object()
 
 
@@ -30,6 +29,14 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         nargs="?",
     )
 
+    group.addoption(
+        "--fzf-bat-preview",
+        action="store_true",
+        dest="use_bat",
+        default=False,
+        help="Use bat as fzf preview command",
+    )
+
 
 def pytest_collection_finish(session: pytest.Session) -> None:
     verbose = session.config.option.verbose
@@ -63,26 +70,25 @@ def pytest_collection_modifyitems(
 
     query = config.option.fzf if config.option.fzf else ""
 
+    def fzf_format(test: pytest.Function) -> str:
+        """Format a test to be displayed with fzf."""
+        assert test.location[1] is not None
+        line_no = test.location[1] + 1
+        return f"{line_no} {test.nodeid}"
+
+    if config.option.use_bat:
+        preview_command = "tail -n +{1} $(echo {2} | cut -d: -f 1)" + f"| {BAT_CMD}"
+    else:
+        preview_command = "pytest-fzf-preview {2}"
+
     kwargs = {
         "multi": True,
         "prompt": "Select test(s): ",
-        # see `fzf_format` for the line format
-        "preview": "tail -n +{1} $(echo {2} | cut -d: -f 1)" + f"| {BAT_CMD}"
-        if BAT_AVAILABLE
-        else "",
+        "preview": preview_command,
         "query": query,
         "cycle": True,
     }
 
-    def fzf_format(test: pytest.Function) -> str:
-        """Format a test to be displayed with fzf.
-
-        This is done to get a proper preview.
-        """
-        assert test.location[1] is not None
-        line_no = test.location[1] + 1
-        return f"{line_no} {test.nodeid}"
-
     res = iterfzf(
         map(fzf_format, items),
         __extra__=[
diff --git a/src/pytest_fzf/previewer.py b/src/pytest_fzf/previewer.py
new file mode 100644
index 0000000..38516fa
--- /dev/null
+++ b/src/pytest_fzf/previewer.py
@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import ast
+import os
+import sys
+
+import pygments
+import pygments.formatters
+import pygments.lexers
+from pygments.styles import get_style_by_name
+
+LEXER = pygments.lexers.PythonLexer()
+
+
+def _get_formatter() -> pygments.formatter.Formatter:
+    style = os.getenv("PYTEST_FZF_THEME") or os.getenv("BAT_THEME")
+
+    if not style:
+        return pygments.formatters.TerminalTrueColorFormatter()
+
+    try:
+        return pygments.formatters.TerminalTrueColorFormatter(
+            style=get_style_by_name(style),
+        )
+    except pygments.util.ClassNotFound:
+        import warnings
+
+        warnings.warn(
+            f'Could not find pygments style: "{style}", falling back to default.',
+            stacklevel=0,
+        )
+        return pygments.formatters.TerminalTrueColorFormatter()
+
+
+def highlight(
+    code: str,
+) -> str:
+    return pygments.highlight(
+        code,
+        lexer=LEXER,
+        formatter=_get_formatter(),
+    )
+
+
+def usage() -> None:
+    command = os.path.basename(sys.argv[0])  # noqa: PTH119
+    print(
+        "Usage: \n"
+        f"  {command} [<nodespec> | <path/to/file.py>"
+        "<object name> [expected_parent_name]]\n"
+        "\nWhere nodespec is a pytest nodespec: path/to/test_module.py::test_name"
+        " or path/to/test_module.py::class_name::test_name"
+        "\n"
+        "Example:\n"
+        "  pytest-fzf-preview tests/test_previewer.py::test_previewer_theme_default\n"
+        "  pytest-fzf-preview tests/test_previewer.py",
+        file=sys.stderr,
+    )
+
+
+def parse_nodespec(nodespec: str) -> tuple[str, str | None, str, list[str]]:
+    """Parse file name and test name from pytest nodespec.
+
+    pytest nodespec: `path/to/test_module.py::class_name::test_name`
+
+    class name is optional
+    """
+    if "[" in nodespec:
+        nodespec, params_str = nodespec.split("[", maxsplit=1)
+
+        params = (
+            params_str[:-1].split(",")  # remove the trailing ]
+            if params_str
+            else []
+        )
+    else:
+        params = []
+
+    file, rest = nodespec.split("::", maxsplit=1)
+    if "::" in rest:
+        class_name, test_name = rest.split("::")
+    else:
+        test_name = rest
+        class_name = None
+
+    return file, class_name, test_name, params
+
+
+def _parse_args(argv: list[str]) -> tuple[str, str | None, str | None]:
+    """Parse file name, and object name from cmdline."""
+    file: str
+    object_name: str | None
+    expected_parent_name: str | None
+
+    if "::" in argv[1]:
+        file, class_name, object_name, _ = parse_nodespec(argv[1])
+
+        expected_parent_name = class_name
+        return file, object_name, expected_parent_name
+
+    file = argv[1]  # path to file
+    object_name = argv[2] if argv[2:] else None
+    expected_parent_name = argv[3] if argv[3:] else None
+    return file, object_name, expected_parent_name
+
+
+def print_highlighted(source: str, start_line: int = 1) -> None:
+    """Print the given source with syntax highlighting and line numbers."""
+    lines = highlight(source).strip().split("\n")
+    max_lineno_width = len(str(len(lines)))
+
+    line_formatter = f"{{lineno:{max_lineno_width}d}}"
+    for index, line in enumerate(lines):
+        line_no = line_formatter.format(lineno=start_line + index)
+        print(f"{line_no} │ {line}")
+
+
+def cli() -> None:
+    if not sys.argv[1:]:
+        usage()
+        sys.exit(1)
+
+    if sys.argv[1] == "-h" or sys.argv[1] == "--help":
+        usage()
+        sys.exit(0)
+
+    file, object_name, parent_name = _parse_args(sys.argv)
+    try:
+        with open(file) as fh:
+            source = fh.read()
+    except FileNotFoundError:
+        print(f"File not found: {file}", file=sys.stderr)
+        sys.exit(1)
+
+    if not object_name:
+        with open(file) as fh:
+            print_highlighted(fh.read())
+        sys.exit(0)
+
+    res = get_source_for_object(
+        source,
+        object_name=object_name,
+    )
+    if not res:
+        print(
+            f"Could not find `{object_name}` in "
+            + (file if file else " the given code")
+            + ".",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    target_source, start_line = res
+    print_highlighted(target_source, start_line)
+
+
+def _parse_tree(
+    tree_or_leaf: ast.AST,
+    object_name: str,
+    expected_parent_name: str | None = None,
+    parent: ast.AST | None = None,
+) -> ast.AST | None:
+    """Parse the given ast for an object with name  "object_name"."""
+    if hasattr(tree_or_leaf, "name") and tree_or_leaf.name == object_name:
+        if not expected_parent_name:
+            return tree_or_leaf
+
+        assert parent is not None
+        return (
+            tree_or_leaf
+            if hasattr(parent, "name") and parent.name == expected_parent_name
+            else None
+        )
+
+    if not hasattr(tree_or_leaf, "body"):
+        return None
+
+    for obj in tree_or_leaf.body:
+        if hasattr(obj, "body") and (
+            res := _parse_tree(
+                obj,
+                object_name,
+                parent=tree_or_leaf,
+                expected_parent_name=expected_parent_name,
+            )
+        ):
+            return res
+
+        if not hasattr(obj, "name"):
+            continue
+
+        if obj.name != object_name:
+            continue
+
+    return None
+
+
+def get_source_for_object(
+    source: str,
+    object_name: str,
+    expected_parent: str | None = None,
+) -> tuple[str, int] | None:
+    """Return the source code of object by parsing the input source."""
+    if res := _parse_tree(
+        ast.parse(source),
+        object_name,
+        expected_parent_name=expected_parent,
+    ):
+        return ast.unparse(res), res.lineno
+
+    return None
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/tests/test_previewer.py b/tests/test_previewer.py
new file mode 100644
index 0000000..886a6b0
--- /dev/null
+++ b/tests/test_previewer.py
@@ -0,0 +1,115 @@
+from contextlib import suppress
+
+import pytest
+
+from pytest_fzf.previewer import _get_formatter, get_source_for_object, parse_nodespec
+
+
+@pytest.mark.parametrize(
+    ("nodespec", "expected"),
+    [
+        (
+            "path/to/test_module.py::className::test_name",
+            ("path/to/test_module.py", "className", "test_name", []),
+        ),
+        (
+            "path/to/test_module.py::test_name",
+            ("path/to/test_module.py", None, "test_name", []),
+        ),
+        (
+            "test_module.py::test_name",
+            ("test_module.py", None, "test_name", []),
+        ),
+        (
+            "test_module.py::test_name[param]",
+            ("test_module.py", None, "test_name", ["param"]),
+        ),
+        (
+            "test_module.py::test_name[param1,param2]",
+            ("test_module.py", None, "test_name", ["param1", "param2"]),
+        ),
+    ],
+)
+def test_parse_nodespec(nodespec, expected):
+    parsed = parse_nodespec(nodespec)
+    # parsed ~= file, name, parent, params
+    assert parsed == expected
+
+
+@pytest.fixture()
+def _clean_theme_env_vars(monkeypatch):
+    """Remove environment vars which modify theme settings."""
+    for env_var in ("BAT_THEME", "PYTEST_FZF_THEME"):
+        with suppress(KeyError):
+            monkeypatch.delenv(env_var)
+
+
+@pytest.mark.parametrize(
+    ("env", "expected"),
+    [
+        (("PYTEST_FZF_THEME", "gruvbox-dark"), "GruvboxDarkStyle"),
+        (("PYTEST_FZF_THEME", "gruvbox-light"), "GruvboxLightStyle"),
+        (("BAT_THEME", "gruvbox-dark"), "GruvboxDarkStyle"),
+        (("BAT_THEME", "gruvbox-light"), "GruvboxLightStyle"),
+    ],
+)
+@pytest.mark.usefixtures("_clean_theme_env_vars")
+def test_previewer_theme(monkeypatch, env, expected):
+    """Test theme overrides using env vars."""
+    monkeypatch.setenv(*env)
+    formatter = _get_formatter()
+    assert formatter.style.__name__ == expected
+
+
+@pytest.mark.usefixtures("_clean_theme_env_vars")
+def test_previewer_theme_default():
+    """Test default highlight theme."""
+    formatter = _get_formatter()
+    assert formatter.style.__name__ == "DefaultStyle"
+
+
+@pytest.mark.usefixtures("_clean_theme_env_vars")
+def test_previewer_theme_invalid(monkeypatch, caplog):
+    """Test default highlight theme."""
+    monkeypatch.setenv("PYTEST_FZF_THEME", "foo")
+
+    with pytest.warns() as warning:
+        formatter = _get_formatter()
+        assert (
+            warning.list[0].message.args[0]
+            == 'Could not find pygments style: "foo", falling back to default.'
+        )
+
+    assert formatter.style.__name__ == "DefaultStyle"
+
+
+def test_get_source_for_object():
+    code = """var = 1
+
+def function(x):
+    print(f"hello {x}")
+
+for el in range(10):
+    function(el)
+"""
+
+    source, line_no = get_source_for_object(code, "function")
+    assert line_no
+    assert source == "def function(x):\n    print(f'hello {x}')"
+
+
+def test_get_source_for_object_class():
+    code = """var = 1
+
+class AnotherClass:
+    def function(self, x):
+        print(f"😡")
+
+class Class:
+    def function(self, x):
+        print(f"hello {x}")
+"""
+
+    source, line_no = get_source_for_object(code, "function", expected_parent="Class")
+    assert line_no
+    assert source == "def function(self, x):\n    print(f'hello {x}')"