From 100a11cd931f557565abd519e335434c3f2a1cb9 Mon Sep 17 00:00:00 2001
From: "joao.nunes" <joao.nunes@bango.com>
Date: Wed, 15 Jan 2025 01:23:21 +0000
Subject: [PATCH] feat(python): support for pypi parse() format and string
 concatenation

---
 src/language/pythonLanguage.ts                | 40 ++++++++++++++++++-
 .../testdata/python/StepDefinitions.py        | 25 +++++++-----
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/src/language/pythonLanguage.ts b/src/language/pythonLanguage.ts
index 88bf1685..a0955122 100644
--- a/src/language/pythonLanguage.ts
+++ b/src/language/pythonLanguage.ts
@@ -34,6 +34,9 @@ export const pythonLanguage: Language = {
     }
   },
   toStepDefinitionExpression(node: TreeSitterSyntaxNode): StringOrRegExp {
+    if (node.type === 'binary_operator') {
+      return collectStringFragments(node).join('')
+    }
     return toStringOrRegExp(node.text)
   },
   defineParameterTypeQueries: [
@@ -95,10 +98,33 @@ export const pythonLanguage: Language = {
       (decorator
         (call
           function: (identifier) @method
-          arguments: (argument_list (string) @expression)
+          arguments: (argument_list [
+            (string) @expression
+            (binary_operator) @expression
+            ])
+        )
+      )
+      (#match? @method "(given|when|then|step)")
+    ) @root`,
+    // pypi parse
+    `(decorator
+      (call
+        function: (identifier) @matcher
+        arguments: (argument_list
+          (call
+            function: [
+              (identifier) @parser
+              (attribute
+                attribute: (identifier) @parser)
+            ]
+            arguments: (argument_list
+              ((_)+ @expression)
+            )
+          )
         )
       )
-      (#match? @method "(given|when|then|step|Given|When|Then|Step)")
+      (#match? @matcher "given|when|then|step")
+      (#match? @parser "parse")
     ) @root`,
   ],
   snippetParameters: {
@@ -157,3 +183,13 @@ export function isRegExp(cleanWord: string): boolean {
 function removePrefix(text: string, prefix: string): string {
   return text.startsWith(prefix) ? text.slice(1) : text
 }
+
+function collectStringFragments(node: TreeSitterSyntaxNode): string[] {
+  if (node.type === 'string') {
+    return [stringLiteral(node.text)]
+  }
+  if (node.type === 'binary_operator') {
+    return node.children.flatMap(collectStringFragments)
+  }
+  return []
+}
diff --git a/test/language/testdata/python/StepDefinitions.py b/test/language/testdata/python/StepDefinitions.py
index d772f217..1ed8cde5 100644
--- a/test/language/testdata/python/StepDefinitions.py
+++ b/test/language/testdata/python/StepDefinitions.py
@@ -1,35 +1,40 @@
-"""Port of givens for testdata."""
-
 from behave import step, given, when, then
+from pytest_bdd import parsers
+from pytest_bdd.parsers import parse
 
 
-@step("a {uuid}")
+@step(parse("a {uuid}"))
 def step_given(context, uuid):
+    """Test PyPi parser syntax"""
     assert uuid
 
 
-@given("a {date}")
+custom_types = {}
+@given(parsers.parse("a {date}", custom_types))
 def step_date(context, date):
+    """Test extra parameter types"""
     assert date
 
 
-@when("a {planet}")
+@when(parse("a " + "{planet}"))
 def step_planet(context, planet):
+    """Test string concatenation"""
     assert planet
 
 
-@then("an {undefined-parameter}")
+@then(parsers.parse("an {undefined-parameter}"))
 def step_undef(context, planet):
+    """Test undefined parameter type"""
     assert planet
 
 
-@Step("/^a regexp$/")
+@when("/^a regexp$/")
 def step_re(context, expression):
-    """Test Re."""
+    """Test regular expression"""
     assert expression
 
 
-@Given("the bee's knees")
+@then("the " + "bee's knees")
 def step_bees(context, expression):
-    """Test Re."""
+    """Test string concatenation"""
     assert expression