From 3d8d7c319dc6f96774539cea97fafd618c85616c Mon Sep 17 00:00:00 2001
From: "Ilya.Kazakevich" <Ilya.Kazakevich@jetbrains.com>
Date: Wed, 25 Jun 2014 19:21:06 +0400
Subject: [PATCH] PY-8427 Behave BDD support  PY-4186 Lettuce BDD tool support

 (create steps added. Need to support several file types)
---
 .../psi/impl/PyCallExpressionHelper.java      |  8 ++---
 .../python/psi/impl/PyFunctionBuilder.java    | 36 +++++++++++++------
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
index b9fe8c342054..74d4b4abaeae 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
@@ -365,14 +365,14 @@ public class PyCallExpressionHelper {
 
   /**
    * Returns argument if it exists and has appropriate type
-   * @param parameter argument
-   * @param argClass expected class
+   * @param parameter  argument
+   * @param argClass   expected class
    * @param expression call expression
-   * @param <T> expected class
+   * @param <T>        expected class
    * @return argument expression or null if has wrong type of does not exist
    */
   @Nullable
-  public static  <T extends PsiElement> T getArgument(
+  public static <T extends PsiElement> T getArgument(
     @NotNull final FunctionParameter parameter,
     @NotNull final Class<T> argClass,
     @NotNull final PyCallExpression expression) {
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java
index 53dea09e4f5c..e0cdf87b2aab 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java
@@ -24,13 +24,9 @@ import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.psi.*;
-
 import org.jetbrains.annotations.NotNull;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 import java.util.regex.Pattern;
 
 /**
@@ -45,10 +41,13 @@ public class PyFunctionBuilder {
   private final List<String> myDecorators = new ArrayList<String>();
   private String myAnnotation = null;
   private String[] myDocStringLines = null;
+  @NotNull
+  private final Map<String, String> myDecoratorValues = new HashMap<String, String>();
 
   /**
    * Creates builder copying signature and doc from another one.
-   * @param source what to copy
+   *
+   * @param source                  what to copy
    * @param decoratorsToCopyIfExist list of decorator names to be copied to new function.
    * @return builder configured by this function
    */
@@ -82,6 +81,7 @@ public class PyFunctionBuilder {
 
   /**
    * Adds docstring to function. Provide doc with out of comment blocks.
+   *
    * @param docString doc
    */
   public void docString(@NotNull final String docString) {
@@ -127,15 +127,21 @@ public class PyFunctionBuilder {
   }
 
   public PyFunction buildFunction(Project project, final LanguageLevel languageLevel) {
-    String text = buildText(project);
     PyElementGenerator generator = PyElementGenerator.getInstance(project);
+    String text = buildText(project, generator, languageLevel);
     return generator.createFromText(languageLevel, PyFunction.class, text);
   }
 
-  private String buildText(Project project) {
+  private String buildText(Project project, PyElementGenerator generator, LanguageLevel languageLevel) {
     StringBuilder builder = new StringBuilder();
     for (String decorator : myDecorators) {
-      builder.append(decorator).append("\n");
+      final StringBuilder decoratorAppender = builder.append('@' + decorator);
+      if (myDecoratorValues.containsKey(decorator)) {
+        final PyCallExpression fakeCall = generator.createCallExpression(languageLevel, "fakeFunction");
+        fakeCall.getArgumentList().addArgument(generator.createStringLiteralFromString(myDecoratorValues.get(decorator)));
+        decoratorAppender.append(fakeCall.getArgumentList().getText());
+      }
+      decoratorAppender.append("\n");
     }
     builder.append("def ");
     builder.append(myName).append("(");
@@ -165,8 +171,18 @@ public class PyFunctionBuilder {
     return builder.toString();
   }
 
+  /**
+   * Adds decorator with argument
+   * @param decoratorName decorator name
+   * @param value its argument
+   */
+  public void decorate(@NotNull final String decoratorName, @NotNull final String value) {
+    decorate(decoratorName);
+    myDecoratorValues.put(decoratorName, value);
+  }
+
   public void decorate(String decoratorName) {
-    myDecorators.add("@" + decoratorName);
+    myDecorators.add(decoratorName);
   }
 
   @NotNull