Get pytest autocompletion in zshell

This question may be better suited for superuser -- if that's the case let me know and I'll shift it.

I use zsh and frequently run pytest from the command line. A very common situation is that I need to run a specific test (or a subtest of a class).

The former looks something like

pytest test/test_foo_file.py::test_foo_function

and the latter something like

pytest test/test_foo_file.py::FooClassTest::test_specific_functionality

It's kind of a pain to write out the entire exact class and test name, and this seems like something ripe for autocompletion or fuzzysearching of some kind. I've been unable to achieve this with what I've found researching -- does anyone have any recommendations?

Let me know if I can be more specific in any way.

Disclaimer: I am not a zsh user, but the approach is pretty much similar to the customizing of bash completions:

1. Create a custom completion file, e.g.

   $ mkdir ~/.zsh-completions
   $ touch ~/.zsh-completions/_pytest
   

2. Inside ~/.zsh-completions/_pytest, write the completion function:

   #compdef pytest
   
   _pytest_complete() {
       local curcontext="$curcontext" state line
       typeset -A opt_args
       compadd "$@" $( pytest --collect-only -q | head -n -2)
   }
   
   _pytest_complete "$@"
   

3. Adjust .zshrc to include custom completions, e.g.

   fpath=(~/.zsh-completions $fpath)
   autoload -U compinit
   compinit
   zstyle ':completion:*' menu select=2
   

Restart the shell. Now you should get the single tests selection on tab completion:

The crucial command here is

$ pytest --collect-only -q | head -n -2

which collects the tests in current directory and lists their names ready to be passed as command line arguments.