fix(python): highlight function calls correctly (#7728)

3 files changed, 77 insertions, 42 deletions

diff --git a/queries/python/highlights.scm b/queries/python/highlights.scm
index d0759e26b..60f3e14c0 100644
--- a/queries/python/highlights.scm
+++ b/queries/python/highlights.scm
@@ -35,48 +35,6 @@
     function: (identifier) @_func))
   (#any-of? @_func "TypeVar" "NewType"))
 
-; Function calls
-(call
-  function: (identifier) @function.call)
-
-(call
-  function: (attribute
-    attribute: (identifier) @function.method.call))
-
-((call
-  function: (identifier) @constructor)
-  (#lua-match? @constructor "^%u"))
-
-((call
-  function: (attribute
-    attribute: (identifier) @constructor))
-  (#lua-match? @constructor "^%u"))
-
-; Decorators
-((decorator
-  "@" @attribute)
-  (#set! priority 101))
-
-(decorator
-  (identifier) @attribute)
-
-(decorator
-  (attribute
-    attribute: (identifier) @attribute))
-
-(decorator
-  (call
-    (identifier) @attribute))
-
-(decorator
-  (call
-    (attribute
-      attribute: (identifier) @attribute)))
-
-((decorator
-  (identifier) @attribute.builtin)
-  (#any-of? @attribute.builtin "classmethod" "property" "staticmethod"))
-
 ; Builtin functions
 ((call
   function: (identifier) @function.builtin)
@@ -450,6 +408,23 @@
       name: (identifier) @constructor)))
   (#any-of? @constructor "__new__" "__init__"))
 
+; Function calls
+(call
+  function: (identifier) @function.call)
+
+(call
+  function: (attribute
+    attribute: (identifier) @function.method.call))
+
+((call
+  function: (identifier) @constructor)
+  (#lua-match? @constructor "^%u"))
+
+((call
+  function: (attribute
+    attribute: (identifier) @constructor))
+  (#lua-match? @constructor "^%u"))
+
 ; Regex from the `re` module
 (call
   function: (attribute
@@ -459,3 +434,28 @@
     (string
       (string_content) @string.regexp))
   (#eq? @_re "re"))
+
+; Decorators
+((decorator
+  "@" @attribute)
+  (#set! priority 101))
+
+(decorator
+  (identifier) @attribute)
+
+(decorator
+  (attribute
+    attribute: (identifier) @attribute))
+
+(decorator
+  (call
+    (identifier) @attribute))
+
+(decorator
+  (call
+    (attribute
+      attribute: (identifier) @attribute)))
+
+((decorator
+  (identifier) @attribute.builtin)
+  (#any-of? @attribute.builtin "classmethod" "property" "staticmethod"))
diff --git a/tests/query/highlights/python/decorators.py b/tests/query/highlights/python/decorators.py
new file mode 100644
index 000000000..278e4e311
--- /dev/null
+++ b/tests/query/highlights/python/decorators.py
@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+
+
+@dataclass
+#^^^^^^^^^ @attribute
+class Data:
+    _foo: str
+
+    @property
+#   ^ @attribute
+#    ^^^^^^^^ @attribute.builtin
+    def foo(self) -> str:
+        return self._foo
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+#^^^^^^ @variable
+#       ^^^^ @variable.member
+#            ^^^^^^^^^^^^^^ @attribute
+def test_func():
+    pass
+
diff --git a/tests/query/highlights/python/functions.py b/tests/query/highlights/python/functions.py
new file mode 100644
index 000000000..6d3028433
--- /dev/null
+++ b/tests/query/highlights/python/functions.py
@@ -0,0 +1,13 @@
+def func() -> None: ...
+
+_ = func()
+#   ^^^^ @function.call
+
+"{}".format(1)
+#    ^^^^^^ @function.method.call
+
+class Foo:
+    def method(self) -> None: ...
+
+Foo().method()
+#     ^^^^^^ @function.method.call