7fc0694fba068ce7f1d01031084803a1d54da96f
[idea/community.git] / python / src / com / jetbrains / python / documentation / docstrings / PyDocstringGenerator.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.documentation.docstrings;
17
18 import com.google.common.base.Preconditions;
19 import com.google.common.collect.Lists;
20 import com.google.common.collect.Maps;
21 import com.google.common.collect.Sets;
22 import com.intellij.codeInsight.CodeInsightUtilCore;
23 import com.intellij.codeInsight.template.*;
24 import com.intellij.openapi.editor.Document;
25 import com.intellij.openapi.editor.Editor;
26 import com.intellij.openapi.fileEditor.FileEditorManager;
27 import com.intellij.openapi.fileEditor.OpenFileDescriptor;
28 import com.intellij.openapi.project.Project;
29 import com.intellij.openapi.util.Condition;
30 import com.intellij.openapi.util.Pair;
31 import com.intellij.openapi.util.TextRange;
32 import com.intellij.openapi.util.text.StringUtil;
33 import com.intellij.openapi.vfs.VirtualFile;
34 import com.intellij.psi.PsiDocumentManager;
35 import com.intellij.psi.PsiElement;
36 import com.intellij.psi.PsiWhiteSpace;
37 import com.intellij.psi.util.PsiTreeUtil;
38 import com.intellij.util.containers.ContainerUtil;
39 import com.jetbrains.python.PyNames;
40 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
41 import com.jetbrains.python.debugger.PySignature;
42 import com.jetbrains.python.debugger.PySignatureCacheManager;
43 import com.jetbrains.python.psi.*;
44 import com.jetbrains.python.toolbox.Substring;
45 import org.jetbrains.annotations.NotNull;
46 import org.jetbrains.annotations.Nullable;
47
48 import java.util.ArrayList;
49 import java.util.List;
50 import java.util.Map;
51 import java.util.Set;
52
53 /**
54  * @author traff
55  */
56 public class PyDocstringGenerator {
57   public static final String TRIPLE_DOUBLE_QUOTES = "\"\"\"";
58   public static final String TRIPLE_SINGLE_QUOTES = "'''";
59   
60   private final List<DocstringParam> myAddedParams = Lists.newArrayList();
61   private final List<DocstringParam> myRemovedParams = Lists.newArrayList();
62   private final String myDocStringText;
63   // Updated after buildAndInsert()
64   @Nullable private PyDocStringOwner myDocStringOwner;
65   private final String myDocStringIndent;
66   private final DocStringFormat myDocStringFormat;
67   private final PsiElement mySettingsAnchor;
68
69   private boolean myUseTypesFromDebuggerSignature = true;
70   private boolean myNewMode = false; // true - generate new string, false - update existing
71   private boolean myAddFirstEmptyLine = false;
72   private boolean myParametersPrepared = false;
73   private String myQuotes = TRIPLE_DOUBLE_QUOTES;
74
75   private PyDocstringGenerator(@Nullable PyDocStringOwner docStringOwner,
76                                @Nullable String docStringText,
77                                @NotNull DocStringFormat format,
78                                @NotNull String indentation,
79                                @NotNull PsiElement settingsAnchor) {
80     myDocStringOwner = docStringOwner;
81     myDocStringIndent = indentation;
82     myDocStringFormat = format;
83     myDocStringText = docStringText;
84     myNewMode = myDocStringText == null;
85     mySettingsAnchor = settingsAnchor;
86   }
87
88   @NotNull
89   public static PyDocstringGenerator forDocStringOwner(@NotNull PyDocStringOwner owner) {
90     String indentation = "";
91     if (owner instanceof PyStatementListContainer) {
92       indentation = PyIndentUtil.getElementIndent(((PyStatementListContainer)owner).getStatementList());
93     }
94     final String docStringText = owner.getDocStringExpression() == null ? null : owner.getDocStringExpression().getText();
95     return new PyDocstringGenerator(owner, docStringText, DocStringUtil.getConfiguredDocStringFormat(owner), indentation, owner);
96   }
97   
98   /**
99    * @param settingsAnchor any PSI element, presumably in the same file/module where generated function is going to be inserted.
100    *                       It's needed to detect configured docstring format and Python indentation size and, as result, 
101    *                       generate properly formatted docstring. 
102    */
103   @NotNull
104   public static PyDocstringGenerator create(@NotNull DocStringFormat format, @NotNull String indentation, @NotNull PsiElement settingsAnchor) {
105     return new PyDocstringGenerator(null, null, format, indentation, settingsAnchor);
106   }
107
108   @NotNull
109   public static PyDocstringGenerator update(@NotNull PyStringLiteralExpression docString) {
110     return new PyDocstringGenerator(PsiTreeUtil.getParentOfType(docString, PyDocStringOwner.class),
111                                     docString.getText(), 
112                                     DocStringUtil.getConfiguredDocStringFormat(docString),
113                                     PyIndentUtil.getElementIndent(docString), 
114                                     docString);
115   }
116
117   /**
118    * @param settingsAnchor any PSI element, presumably in the same file/module where generated function is going to be inserted.
119    *                       It's needed to detect configured docstring format and Python indentation size and, as result, 
120    *                       generate properly formatted docstring. 
121    */
122   @NotNull
123   public static PyDocstringGenerator update(@NotNull DocStringFormat format,
124                                             @NotNull String indentation,
125                                             @NotNull String text, PsiElement settingsAnchor) {
126     return new PyDocstringGenerator(null, text, format, indentation, settingsAnchor);
127   }
128
129   @NotNull
130   public PyDocstringGenerator withParam(@NotNull String name) {
131     return withParamTypedByName(name, null);
132   }
133
134   @NotNull
135   public PyDocstringGenerator withParam(@NotNull PyNamedParameter param) {
136     return withParam(getPreferredParameterName(param));
137   }
138
139   @NotNull
140   public PyDocstringGenerator withParamTypedByName(@NotNull String name, @Nullable String type) {
141     myAddedParams.add(new DocstringParam(name, type, false));
142     return this;
143   }
144
145   @NotNull
146   public PyDocstringGenerator withParamTypedByName(@NotNull PyNamedParameter name, @Nullable String type) {
147     return withParamTypedByName(getPreferredParameterName(name), type);
148   }
149
150   @NotNull
151   public PyDocstringGenerator withReturnValue(@Nullable String type) {
152     myAddedParams.add(new DocstringParam("", type, true));
153     return this;
154   }
155
156   @NotNull
157   public PyDocstringGenerator withoutParam(@NotNull String name) {
158     myRemovedParams.add(new DocstringParam(name, null, false));
159     return this;
160   }
161
162   @NotNull
163   public PyDocstringGenerator withQuotes(@NotNull String quotes) {
164     myQuotes = quotes;
165     return this;
166   }
167
168   @NotNull
169   public PyDocstringGenerator useTypesFromDebuggerSignature(boolean use) {
170     myUseTypesFromDebuggerSignature = use;
171     return this;
172   }
173
174   @NotNull
175   public PyDocstringGenerator addFirstEmptyLine() {
176     myAddFirstEmptyLine = true;
177     return this;
178   }
179
180
181   @NotNull
182   public PyDocstringGenerator forceNewMode() {
183     myNewMode = true;
184     return this;
185   }
186
187   /**
188    * @param addReturn by default return declaration is added only if function body contains return statement. Sometimes it's not
189    *                  possible, e.g. in  {@link com.jetbrains.python.editor.PythonEnterHandler} where unclosed docstring literal
190    *                  "captures" whole function body including return statements. Keep in mind that declaration for the return value
191    *                  won't be added if containing function is <tt>__init__</tt> or <tt>__new__</tt> method.
192    */
193   @NotNull
194   public PyDocstringGenerator withInferredParameters(boolean addReturn) {
195     if (myDocStringOwner instanceof PyFunction) {
196       for (PyParameter param : ((PyFunction)myDocStringOwner).getParameterList().getParameters()) {
197         if (param.getAsNamed() == null) {
198           continue;
199         }
200         final String paramName = param.getName();
201         final StructuredDocString docString = getStructuredDocString();
202         if (StringUtil.isEmpty(paramName) || param.isSelf() || docString != null && docString.getParameters().contains(paramName)) {
203           continue;
204         }
205         withParam((PyNamedParameter)param);
206       }
207       final RaiseVisitor visitor = new RaiseVisitor();
208       final PyStatementList statementList = ((PyFunction)myDocStringOwner).getStatementList();
209       statementList.accept(visitor);
210       if (!isConstructor((PyFunction)myDocStringOwner) && (visitor.myHasReturn || addReturn)) {
211         // will add :return: placeholder in Sphinx/Epydoc docstrings
212         withReturnValue(null);
213       }
214     }
215     return this;
216   }
217
218   private static boolean isConstructor(@NotNull PyFunction function) {
219     final String funcName = function.getName();
220     return PyNames.INIT.equals(funcName) && function.getContainingClass() != null;
221   }
222
223   @NotNull
224   public String getDocStringIndent() {
225     return myDocStringIndent;
226   }
227
228   @NotNull
229   public DocStringFormat getDocStringFormat() {
230     return myDocStringFormat;
231   }
232
233   public boolean isNewMode() {
234     return myNewMode;
235   }
236
237   /**
238    * Populate parameters for function if nothing was specified.
239    * Order parameters, remove duplicates and merge parameters with and without type according to docstring format.
240    */
241   private void prepareParameters() {
242     if (myParametersPrepared) {
243       return;
244     }
245     final Set<Pair<String, Boolean>> withoutType = Sets.newHashSet();
246     final Map<Pair<String, Boolean>, String> paramTypes = Maps.newHashMap();
247     for (DocstringParam param : myAddedParams) {
248       if (param.getType() == null) {
249         withoutType.add(Pair.create(param.getName(), param.isReturnValue()));
250       }
251       else {
252         // leave only the last type for parameter
253         paramTypes.put(Pair.create(param.getName(), param.isReturnValue()), param.getType());
254       }
255     }
256
257     // Sanitize parameters
258     PySignature signature = null;
259     if (myDocStringOwner instanceof PyFunction && myUseTypesFromDebuggerSignature) {
260       signature = PySignatureCacheManager.getInstance(myDocStringOwner.getProject()).findSignature((PyFunction)myDocStringOwner);
261     }
262     final DocStringFormat format = myDocStringFormat;
263     final ArrayList<DocstringParam> filtered = Lists.newArrayList();
264     final Set<Pair<String, Boolean>> processed = Sets.newHashSet();
265     for (DocstringParam param : myAddedParams) {
266       final Pair<String, Boolean> paramCoordinates = Pair.create(param.getName(), param.isReturnValue());
267       if (processed.contains(paramCoordinates)) {
268         continue;
269       }
270       if (param.getType() == null) {
271         String type = paramTypes.get(paramCoordinates);
272         if (type == null && PyCodeInsightSettings.getInstance().INSERT_TYPE_DOCSTUB) {
273           if (signature != null) {
274             type = StringUtil.notNullize(param.isReturnValue() ? 
275                                          signature.getReturnTypeQualifiedName() :
276                                          signature.getArgTypeQualifiedName(param.getName()));
277           }
278           else {
279             type = "";
280           }
281         }
282         if (type != null) {
283           // Google and Numpy docstring formats combine type and description in single declaration, thus
284           // if both declaration with type and without it are requested, we should filter out duplicates
285           if (format == DocStringFormat.GOOGLE || format == DocStringFormat.NUMPY) {
286             filtered.add(new DocstringParam(param.getName(), type, param.isReturnValue()));
287           }
288           else {
289             // In reST and Epydoc for each parameter add two tags, e.g. in reST (Sphinx)
290             // :param foo:
291             // :type foo:
292             filtered.add(param);
293             filtered.add(new DocstringParam(param.getName(), type, param.isReturnValue()));
294           }
295         }
296         else {
297           // no type was given and it's not required by settings
298           filtered.add(param);
299         }
300       }
301       else if (!withoutType.contains(paramCoordinates)) {
302         filtered.add(param);
303       }
304       processed.add(paramCoordinates);
305     }
306     myAddedParams.clear();
307     myAddedParams.addAll(filtered);
308     myParametersPrepared = true;
309   }
310
311   public boolean hasParametersToAdd() {
312     prepareParameters();
313     return !myAddedParams.isEmpty();
314   }
315
316   @Nullable
317   private PyStringLiteralExpression getDocStringExpression() {
318     Preconditions.checkNotNull(myDocStringOwner, "For this action docstring owner must be supplied");
319     return myDocStringOwner.getDocStringExpression();
320   }
321
322   @Nullable
323   private StructuredDocString getStructuredDocString() {
324     return myDocStringText == null ? null : DocStringUtil.parseDocString(myDocStringFormat, myDocStringText);
325   }
326
327   public void startTemplate() {
328     Preconditions.checkNotNull(myDocStringOwner, "For this action docstring owner must be supplied");
329     final PyStringLiteralExpression docStringExpression = getDocStringExpression();
330     assert docStringExpression != null;
331
332     final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression);
333
334     if (myAddedParams.size() > 1) {
335       throw new IllegalArgumentException("TemplateBuilder can be created only for one parameter");
336     }
337
338     final DocstringParam paramToEdit = getParamToEdit();
339     final DocStringFormat format = myDocStringFormat;
340     if (format == DocStringFormat.PLAIN) {
341       return;
342     }
343     final StructuredDocString parsed = DocStringUtil.parseDocString(format, docStringExpression);
344     final Substring substring;
345     if (paramToEdit.isReturnValue()) {
346       substring = parsed.getReturnTypeSubstring();
347     }
348     else {
349       final String paramName = paramToEdit.getName();
350       substring = parsed.getParamTypeSubstring(paramName);
351     }
352     if (substring == null) {
353       return;
354     }
355     builder.replaceRange(substring.getTextRange(), getDefaultType(getParamToEdit()));
356     Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
357     final VirtualFile virtualFile = myDocStringOwner.getContainingFile().getVirtualFile();
358     if (virtualFile == null) return;
359     final Project project = myDocStringOwner.getProject();
360     OpenFileDescriptor descriptor = new OpenFileDescriptor(project, virtualFile, docStringExpression.getTextOffset());
361     Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
362     if (targetEditor != null) {
363       TemplateManager.getInstance(project).startTemplate(targetEditor, template);
364     }
365   }
366
367   @NotNull
368   public String getPreferredParameterName(@NotNull PyNamedParameter parameter) {
369     if (getDocStringFormat() == DocStringFormat.GOOGLE && parameter.getAsNamed() != null) {
370       return parameter.getAsNamed().getRepr(false);
371     }
372     return StringUtil.notNullize(parameter.getName());
373   }
374
375
376   @NotNull
377   private static String getDefaultType(@NotNull DocstringParam param) {
378     if (StringUtil.isEmpty(param.getType())) {
379       return PyNames.OBJECT;
380     }
381     else {
382       return param.getType();
383     }
384   }
385
386   @NotNull
387   public String buildDocString() {
388     prepareParameters();
389     if (myNewMode) {
390       return createDocString();
391     }
392     else {
393       return updateDocString();
394     }
395   }
396
397   @NotNull
398   private String createDocString() {
399     DocStringBuilder builder = null;
400     if (myDocStringFormat == DocStringFormat.EPYTEXT || myDocStringFormat == DocStringFormat.REST) {
401       builder = new TagBasedDocStringBuilder(myDocStringFormat == DocStringFormat.EPYTEXT ? "@" : ":");
402       TagBasedDocStringBuilder tagBuilder = (TagBasedDocStringBuilder)builder;
403       if (myAddFirstEmptyLine) {
404         tagBuilder.addEmptyLine();
405       }
406       for (DocstringParam param : myAddedParams) {
407         if (param.isReturnValue()) {
408           if (param.getType() != null) {
409             tagBuilder.addReturnValueType(param.getType());
410           }
411           else {
412             tagBuilder.addReturnValueDescription("");
413           }
414         }
415         else {
416           if (param.getType() != null) {
417             tagBuilder.addParameterType(param.getName(), param.getType());
418           }
419           else {
420             tagBuilder.addParameterDescription(param.getName(), "");
421           }
422         }
423       }
424     }
425     else if (myDocStringFormat == DocStringFormat.GOOGLE || myDocStringFormat == DocStringFormat.NUMPY) {
426       builder = myDocStringFormat == DocStringFormat.GOOGLE ? GoogleCodeStyleDocStringBuilder.forProject(mySettingsAnchor.getProject()) : new NumpyDocStringBuilder();
427       final SectionBasedDocStringBuilder sectionBuilder = (SectionBasedDocStringBuilder)builder;
428       if (myAddFirstEmptyLine) {
429         sectionBuilder.addEmptyLine();
430       }
431       final List<DocstringParam> parameters = ContainerUtil.findAll(myAddedParams, param -> !param.isReturnValue());
432       if (!parameters.isEmpty()) {
433         sectionBuilder.startParametersSection();
434         for (DocstringParam param : parameters) {
435           sectionBuilder.addParameter(param.getName(), param.getType(), "");
436         }
437       }
438
439       final List<DocstringParam> returnValues = ContainerUtil.findAll(myAddedParams, param -> param.isReturnValue());
440
441       if (!returnValues.isEmpty()) {
442         sectionBuilder.startReturnsSection();
443         boolean hasTypedReturns = false;
444         for (DocstringParam returnValue : returnValues) {
445           if (StringUtil.isNotEmpty(returnValue.getType())) {
446             sectionBuilder.addReturnValue(null, returnValue.getType(), "");
447             hasTypedReturns = true;
448           }
449         }
450         if (!hasTypedReturns) {
451           sectionBuilder.addEmptyLine();
452         }
453       }
454     }
455     if (builder != null && !builder.getLines().isEmpty()) {
456       return myQuotes + '\n' + builder.buildContent(myDocStringIndent, true) + '\n' + myDocStringIndent + myQuotes;
457     }
458     return createEmptyFallbackDocString();
459   }
460
461   @NotNull
462   private String updateDocString() {
463     DocStringUpdater updater = null;
464     if (myDocStringFormat == DocStringFormat.EPYTEXT || myDocStringFormat == DocStringFormat.REST) {
465       final String prefix = myDocStringFormat == DocStringFormat.EPYTEXT ? "@" : ":";
466       //noinspection unchecked,ConstantConditions
467       updater = new TagBasedDocStringUpdater((TagBasedDocString)getStructuredDocString(), prefix, myDocStringIndent);
468     }
469     else if (myDocStringFormat == DocStringFormat.GOOGLE) {
470       //noinspection ConstantConditions
471       updater = GoogleCodeStyleDocStringUpdater.forProject((GoogleCodeStyleDocString)getStructuredDocString(), 
472                                                            myDocStringIndent, 
473                                                            mySettingsAnchor.getProject());
474     }
475     else if (myDocStringFormat == DocStringFormat.NUMPY) {
476       //noinspection ConstantConditions
477       updater = new NumpyDocStringUpdater((SectionBasedDocString)getStructuredDocString(), myDocStringIndent);
478     }
479     // plain docstring - do nothing
480     else if (myDocStringText != null){
481       return myDocStringText;
482     }
483     if (updater != null) {
484       for (DocstringParam param : myAddedParams) {
485         if (param.isReturnValue()) {
486           updater.addReturnValue(param.getType());
487         }
488         else {
489           updater.addParameter(param.getName(), param.getType());
490         }
491       }
492       for (DocstringParam param : myRemovedParams) {
493         if (!param.isReturnValue()) {
494           updater.removeParameter(param.getName());
495         }
496       }
497       return updater.getDocStringText();
498     }
499     return createEmptyFallbackDocString();
500   }
501
502   @NotNull
503   private String createEmptyFallbackDocString() {
504     return myQuotes + '\n' + myDocStringIndent + myQuotes;
505   }
506
507   private DocstringParam getParamToEdit() {
508     if (myAddedParams.size() == 0) {
509       throw new IllegalStateException("We should have at least one param to edit");
510     }
511     return myAddedParams.get(0);
512   }
513
514   @NotNull
515   public PyDocStringOwner buildAndInsert() {
516     Preconditions.checkNotNull(myDocStringOwner, "For this action docstring owner must be supplied");
517     final String replacementText = buildDocString();
518
519     final Project project = myDocStringOwner.getProject();
520     PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
521     final PyExpressionStatement replacement = elementGenerator.createDocstring(replacementText);
522
523     final PyStringLiteralExpression docStringExpression = getDocStringExpression();
524     if (docStringExpression != null) {
525       docStringExpression.replace(replacement.getExpression());
526     }
527     else {
528       PyStatementListContainer container = PyUtil.as(myDocStringOwner, PyStatementListContainer.class);
529       if (container == null) {
530         throw new IllegalStateException("Should be a function or class");
531       }
532       final PyStatementList statements = container.getStatementList();
533       final String indentation = PyIndentUtil.getElementIndent(statements);
534       final Document document = PsiDocumentManager.getInstance(project).getDocument(myDocStringOwner.getContainingFile());
535
536       if (document != null) {
537         final PsiElement beforeStatements = statements.getPrevSibling();
538         final boolean onSameLine = !(beforeStatements instanceof PsiWhiteSpace) || !beforeStatements.textContains('\n');
539         if (onSameLine || statements.getStatements().length == 0) {
540           String replacementWithLineBreaks = "\n" + indentation + replacementText;
541           if (statements.getStatements().length > 0) {
542             replacementWithLineBreaks += "\n" + indentation;
543           }
544           final PsiDocumentManager documentManager = PsiDocumentManager.getInstance(project);
545           documentManager.doPostponedOperationsAndUnblockDocument(document);
546           final TextRange range = beforeStatements.getTextRange();
547           try {
548             if (beforeStatements instanceof PsiWhiteSpace) {
549               if (statements.getStatements().length > 0) {
550                 document.replaceString(range.getStartOffset(), range.getEndOffset(), replacementWithLineBreaks);
551               }
552               else {
553                 // preserve original spacing, since it probably separates function and other declarations
554                 document.insertString(range.getStartOffset(), replacementWithLineBreaks);
555               }
556             }
557             else {
558               document.insertString(range.getEndOffset(), replacementWithLineBreaks);
559             }
560           }
561           finally {
562             documentManager.commitDocument(document);
563           }
564         }
565         else {
566           statements.addBefore(replacement, statements.getStatements()[0]);
567         }
568       }
569     }
570     myDocStringOwner = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(myDocStringOwner);
571     return myDocStringOwner;
572   }
573
574   public static class DocstringParam {
575
576     private final String myName;
577     private final String myType;
578     private final boolean myReturnValue;
579     private DocstringParam(@NotNull String name, @Nullable String type, boolean isReturn) {
580       myName = name;
581       myType = type;
582       myReturnValue = isReturn;
583     }
584
585     @NotNull
586     public String getName() {
587       return myName;
588     }
589
590     @Nullable
591     public String getType() {
592       return myType;
593     }
594
595     public boolean isReturnValue() {
596       return myReturnValue;
597     }
598
599     @Override
600     public boolean equals(Object o) {
601       if (this == o) return true;
602       if (o == null || getClass() != o.getClass()) return false;
603
604       DocstringParam param = (DocstringParam)o;
605
606       if (myReturnValue != param.myReturnValue) return false;
607       if (!myName.equals(param.myName)) return false;
608       if (myType != null ? !myType.equals(param.myType) : param.myType != null) return false;
609
610       return true;
611     }
612
613     @Override
614     public int hashCode() {
615       int result = myName.hashCode();
616       result = 31 * result + (myType != null ? myType.hashCode() : 0);
617       result = 31 * result + (myReturnValue ? 1 : 0);
618       return result;
619     }
620
621     @Override
622     public String toString() {
623       return "DocstringParam{" +
624              "myName='" + myName + '\'' +
625              ", myType='" + myType + '\'' +
626              ", myReturnValue=" + myReturnValue +
627              '}';
628     }
629
630   }
631   private static class RaiseVisitor extends PyRecursiveElementVisitor {
632
633     private boolean myHasRaise = false;
634     private boolean myHasReturn = false;
635     @Nullable private PyExpression myRaiseTarget = null;
636     @Override
637     public void visitPyRaiseStatement(@NotNull PyRaiseStatement node) {
638       myHasRaise = true;
639       final PyExpression[] expressions = node.getExpressions();
640       if (expressions.length > 0) {
641         myRaiseTarget = expressions[0];
642       }
643     }
644
645     @Override
646     public void visitPyReturnStatement(PyReturnStatement node) {
647       myHasReturn = true;
648     }
649
650     @NotNull
651     public String getRaiseTargetText() {
652       if (myRaiseTarget != null) {
653         String raiseTarget = myRaiseTarget.getText();
654         if (myRaiseTarget instanceof PyCallExpression) {
655           final PyExpression callee = ((PyCallExpression)myRaiseTarget).getCallee();
656           if (callee != null) {
657             raiseTarget = callee.getText();
658           }
659         }
660         return raiseTarget;
661       }
662       return "";
663     }
664
665   }
666 }
667