diff --git a/queries/python/context.scm b/queries/python/context.scm
index 1b1368e0..0c70479b 100644
--- a/queries/python/context.scm
+++ b/queries/python/context.scm
@@ -1,6 +1,14 @@
 (class_definition
   body: (_) @context.end) @context
 
+; If the function starts with a comment, then use that to end the context.
+;
+; Note: priority of overlapping rules isn't defined, but this one seems to
+; always take precedence over the one below, regardless of the order in which
+; they are defined. Maybe the smallest match wins?
+(function_definition
+  (comment) @context.end) @context
+
 (function_definition
   body: (_) @context.end) @context
 
diff --git a/test/lang/test.py b/test/lang/test.py
index be398f14..81e4cc74 100644
--- a/test/lang/test.py
+++ b/test/lang/test.py
@@ -1,6 +1,18 @@
+# {{TEST}}
 def hello(name: str, age: int) -> None: # {{CONTEXT}}
     print(f"Hello {name}! You are {age} years old.")
 
 
 
     # {{CURSOR}}
+
+# {{TEST}}
+def hello2(                # {{CONTEXT}}
+        name: str,         # {{CONTEXT}}
+        age: int) -> None: # {{CONTEXT}}
+    # comment
+
+
+
+
+    name = 'barry' # {{CURSOR}}