1; subclasses of unittest.TestCase or TestCase
  2(
  3    (class_definition
  4        name: (identifier) @run @_unittest_class_name
  5        superclasses: (argument_list
  6            [(identifier) @_superclass
  7            (attribute (identifier) @_superclass)]
  8        )
  9        (#eq? @_superclass "TestCase")
 10    ) @_python-unittest-class
 11    (#set! tag python-unittest-class)
 12)
 13
 14; test methods whose names start with `test` in a TestCase
 15(
 16    (class_definition
 17        name: (identifier) @_unittest_class_name
 18        superclasses: (argument_list
 19            [(identifier) @_superclass
 20            (attribute (identifier) @_superclass)]
 21        )
 22        (#eq? @_superclass "TestCase")
 23        body: (block
 24                (function_definition
 25                    name: (identifier) @run @_unittest_method_name
 26                    (#match? @_unittest_method_name "^test.*")
 27                ) @_python-unittest-method
 28                (#set! tag python-unittest-method)
 29            )
 30        )
 31)
 32
 33; pytest functions
 34(
 35    (module
 36        (function_definition
 37            name: (identifier) @run @_pytest_method_name
 38            (#match? @_pytest_method_name "^test_")
 39        ) @_python-pytest-method
 40    )
 41    (#set! tag python-pytest-method)
 42)
 43
 44; decorated pytest functions
 45(
 46    (module
 47        (decorated_definition
 48            (decorator)+ @_decorator
 49            definition: (function_definition
 50                name: (identifier) @run @_pytest_method_name
 51                (#match? @_pytest_method_name "^test_")
 52                )
 53            ) @_python-pytest-method
 54        )
 55    (#set! tag python-pytest-method)
 56)
 57
 58; pytest classes
 59(
 60    (module
 61        (class_definition
 62            name: (identifier) @run @_pytest_class_name
 63            (#match? @_pytest_class_name "^Test")
 64        )
 65        (#set! tag python-pytest-class)
 66    )
 67)
 68
 69; pytest class methods
 70(
 71    (module
 72        (class_definition
 73            name: (identifier) @_pytest_class_name
 74            (#match? @_pytest_class_name "^Test")
 75            body: (block
 76                    (function_definition
 77                        name: (identifier) @run @_pytest_method_name
 78                        (#match? @_pytest_method_name "^test")
 79                    ) @_python-pytest-method
 80                    (#set! tag python-pytest-method)
 81            )
 82        )
 83    )
 84)
 85
 86; decorated pytest class methods
 87(
 88    (module
 89        (class_definition
 90            name: (identifier) @_pytest_class_name
 91            (#match? @_pytest_class_name "^Test")
 92            body: (block
 93                (decorated_definition
 94                    (decorator)+ @_decorator
 95                    definition: (function_definition
 96                        name: (identifier) @run @_pytest_method_name
 97                        (#match? @_pytest_method_name "^test_")
 98                        )
 99                    )
100            ) @_python-pytest-method
101            (#set! tag python-pytest-method)
102        )
103    )
104)
105
106; module main method
107(
108    (module
109        (if_statement
110            condition: (comparison_operator
111                (identifier) @run @_lhs
112                operators: "=="
113                (string) @_rhs
114            )
115            (#eq? @_lhs "__name__")
116            (#match? @_rhs "^[\"']__main__[\"']$")
117            (#set! tag python-module-main-method)
118        )
119    )
120)