e2456210d26841bfc5ee0104f4b7a57127730773
[idea/community.git] / python / src / com / jetbrains / python / psi / impl / PyFunctionBuilder.java
1 /*
2  * Copyright 2000-2014 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.psi.impl;
17
18 import com.intellij.openapi.project.Project;
19 import com.intellij.openapi.util.text.StringUtil;
20 import com.intellij.psi.PsiElement;
21 import com.intellij.util.ArrayUtil;
22 import com.jetbrains.python.PyNames;
23 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
24 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
25 import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator;
26 import com.jetbrains.python.psi.*;
27 import org.jetbrains.annotations.NotNull;
28
29 import java.util.*;
30
31 /**
32  * @author yole
33  */
34 public class PyFunctionBuilder {
35   private final String myName;
36   private final List<String> myParameters = new ArrayList<>();
37   private final List<String> myStatements = new ArrayList<>();
38   private final List<String> myDecorators = new ArrayList<>();
39   private String myAnnotation = null;
40   @NotNull
41   private final Map<String, String> myDecoratorValues = new HashMap<>();
42   private boolean myAsync = false;
43   private PyDocstringGenerator myDocStringGenerator;
44
45   /**
46    * Creates builder copying signature and doc from another one.
47    *
48    * @param source                  what to copy
49    * @param decoratorsToCopyIfExist list of decorator names to be copied to new function.
50    * @return builder configured by this function
51    */
52   @NotNull
53   public static PyFunctionBuilder copySignature(@NotNull final PyFunction source, @NotNull final String... decoratorsToCopyIfExist) {
54     final String name = source.getName();
55     final PyFunctionBuilder functionBuilder = new PyFunctionBuilder((name != null) ? name : "", source);
56     for (final PyParameter parameter : source.getParameterList().getParameters()) {
57       final String parameterName = parameter.getName();
58       if (parameterName != null) {
59         functionBuilder.parameter(parameterName);
60       }
61     }
62     final PyDecoratorList decoratorList = source.getDecoratorList();
63     if (decoratorList != null) {
64       for (final PyDecorator decorator : decoratorList.getDecorators()) {
65         final String decoratorName = decorator.getName();
66         if (decoratorName != null) {
67           if (ArrayUtil.contains(decoratorName, decoratorsToCopyIfExist)) {
68             functionBuilder.decorate(decoratorName);
69           }
70         }
71       }
72     }
73     functionBuilder.myDocStringGenerator = PyDocstringGenerator.forDocStringOwner(source);
74     return functionBuilder;
75   }
76
77   @Deprecated
78   public PyFunctionBuilder(@NotNull String name) {
79     myName = name;
80     myDocStringGenerator = null;
81   }
82
83   /**
84    * @param settingsAnchor any PSI element, presumably in the same file/module where generated function is going to be inserted.
85    *                       It's needed to detect configured docstring format and Python indentation size and, as result, 
86    *                       generate properly formatted docstring. 
87    */
88   public PyFunctionBuilder(@NotNull String name, @NotNull PsiElement settingsAnchor) {
89     myName = name;
90     myDocStringGenerator = PyDocstringGenerator.create(DocStringUtil.getConfiguredDocStringFormat(settingsAnchor), 
91                                                        PyIndentUtil.getIndentFromSettings(settingsAnchor.getProject()), 
92                                                        settingsAnchor);
93   }
94
95   /**
96    * Adds param and its type to doc
97    * @param format what docstyle to use to doc param type
98    * @param name param name
99    * @param type param type
100    */
101   @NotNull
102   public PyFunctionBuilder parameterWithType(@NotNull String name, @NotNull String type) {
103     parameter(name);
104     myDocStringGenerator.withParamTypedByName(name, type);
105     return this;
106   }
107
108   @NotNull
109   @Deprecated
110   public PyFunctionBuilder parameterWithType(@NotNull final String name,
111                                              @NotNull final String type,
112                                              @NotNull final DocStringFormat format) {
113     parameter(name);
114     myDocStringGenerator.withParamTypedByName(name, type);
115     return this;
116   }
117
118   public PyFunctionBuilder parameter(String baseName) {
119     String name = baseName;
120     int uniqueIndex = 0;
121     while (myParameters.contains(name)) {
122       uniqueIndex++;
123       name = baseName + uniqueIndex;
124     }
125     myParameters.add(name);
126     return this;
127   }
128
129   public PyFunctionBuilder annotation(String text) {
130     myAnnotation = text;
131     return this;
132   }
133
134   public PyFunctionBuilder makeAsync() {
135     myAsync = true;
136     return this;
137   }
138
139   public PyFunctionBuilder statement(String text) {
140     myStatements.add(text);
141     return this;
142   }
143
144   public PyFunction addFunction(PsiElement target, final LanguageLevel languageLevel) {
145     return (PyFunction)target.add(buildFunction(target.getProject(), languageLevel));
146   }
147
148   public PyFunction addFunctionAfter(PsiElement target, PsiElement anchor, final LanguageLevel languageLevel) {
149     return (PyFunction)target.addAfter(buildFunction(target.getProject(), languageLevel), anchor);
150   }
151
152   public PyFunction buildFunction(Project project, final LanguageLevel languageLevel) {
153     PyElementGenerator generator = PyElementGenerator.getInstance(project);
154     String text = buildText(project, generator, languageLevel);
155     return generator.createFromText(languageLevel, PyFunction.class, text);
156   }
157
158   private String buildText(Project project, PyElementGenerator generator, LanguageLevel languageLevel) {
159     StringBuilder builder = new StringBuilder();
160     for (String decorator : myDecorators) {
161       final StringBuilder decoratorAppender = builder.append('@' + decorator);
162       if (myDecoratorValues.containsKey(decorator)) {
163         final PyCallExpression fakeCall = generator.createCallExpression(languageLevel, "fakeFunction");
164         fakeCall.getArgumentList().addArgument(generator.createStringLiteralFromString(myDecoratorValues.get(decorator)));
165         decoratorAppender.append(fakeCall.getArgumentList().getText());
166       }
167       decoratorAppender.append("\n");
168     }
169     if (myAsync) {
170       builder.append("async ");
171     }
172     builder.append("def ");
173     builder.append(myName).append("(");
174     builder.append(StringUtil.join(myParameters, ", "));
175     builder.append(")");
176     if (myAnnotation != null) {
177       builder.append(myAnnotation);
178     }
179     builder.append(":");
180     List<String> statements = myStatements.isEmpty() ? Collections.singletonList(PyNames.PASS) : myStatements;
181
182     final String indent = PyIndentUtil.getIndentFromSettings(project);
183     // There was original docstring or some parameters were added via parameterWithType()
184     if (!myDocStringGenerator.isNewMode() || myDocStringGenerator.hasParametersToAdd()) {
185       final String docstring = PyIndentUtil.changeIndent(myDocStringGenerator.buildDocString(), true, indent);
186       builder.append('\n').append(indent).append(docstring);
187     }
188     for (String statement : statements) {
189       builder.append('\n').append(indent).append(statement);
190     }
191     return builder.toString();
192   }
193
194   /**
195    * Adds decorator with argument
196    *
197    * @param decoratorName decorator name
198    * @param value         its argument
199    */
200   public void decorate(@NotNull final String decoratorName, @NotNull final String value) {
201     decorate(decoratorName);
202     myDecoratorValues.put(decoratorName, value);
203   }
204
205   public void decorate(String decoratorName) {
206     myDecorators.add(decoratorName);
207   }
208 }