diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java index 344185608f50..ba0d519f8d4b 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java @@ -20,6 +20,7 @@ import com.intellij.openapi.util.text.StringUtil; import com.intellij.psi.PsiElement; import com.intellij.util.ArrayUtil; import com.jetbrains.python.PyNames; +import com.jetbrains.python.documentation.docstrings.DocStringFormat; import com.jetbrains.python.documentation.docstrings.DocStringUtil; import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator; import com.jetbrains.python.psi.*; @@ -73,6 +74,11 @@ public class PyFunctionBuilder { return functionBuilder; } + public PyFunctionBuilder(@NotNull String name) { + myName = name; + myDocStringGenerator = null; + } + public PyFunctionBuilder(@NotNull String name, @NotNull PsiElement settingsAnchor) { myName = name; myDocStringGenerator = PyDocstringGenerator.create(DocStringUtil.getConfiguredDocStringFormat(settingsAnchor), @@ -93,6 +99,15 @@ public class PyFunctionBuilder { return this; } + @NotNull + public PyFunctionBuilder parameterWithType(@NotNull final String name, + @NotNull final String type, + @NotNull final DocStringFormat format) { + parameter(name); + myDocStringGenerator.withParamTypedByName(name, type); + return this; + } + public PyFunctionBuilder parameter(String baseName) { String name = baseName; int uniqueIndex = 0;