TitleCapitalizationInspection: multiple returns

plugins/devkit/src/inspections/TitleCapitalizationInspection.java

@@ -28,6 +28,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.ReadonlyStatusHandler;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
+import com.intellij.psi.util.InheritanceUtil;
 import com.intellij.psi.util.PropertyUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtil;
@@ -36,6 +37,7 @@ import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Collection;
 import java.util.Collections;
 
 /**
@@ -74,14 +76,15 @@ public class TitleCapitalizationInspection extends BaseJavaLocalInspectionTool {
       @Override
       public void visitMethod(PsiMethod method) {
         PsiType type = method.getReturnType();
-        if (!(type instanceof PsiClassType) || !"String".equals(((PsiClassType)type).getClassName())) return;
-        PsiReturnStatement returnStatement = PsiTreeUtil.findChildOfType(method, PsiReturnStatement.class);
-        if (returnStatement == null) return;
-        PsiExpression expression = returnStatement.getReturnValue();
-        String value = getTitleValue(expression);
-        if (value == null) return;
-        Nls.Capitalization capitalization = getCapitalizationFromAnno(method);
-        checkCapitalization(expression, holder, capitalization);
+        if (!InheritanceUtil.isInheritor(type, CommonClassNames.JAVA_LANG_STRING)) return;
+        Collection<PsiReturnStatement> statements = PsiTreeUtil.findChildrenOfType(method, PsiReturnStatement.class);
+        for (PsiReturnStatement returnStatement : statements) {
+          PsiExpression expression = returnStatement.getReturnValue();
+          String value = getTitleValue(expression);
+          if (value == null) continue;
+          Nls.Capitalization capitalization = getCapitalizationFromAnno(method);
+          checkCapitalization(expression, holder, capitalization);
+        }
       }
 
       @Override

plugins/devkit/testData/inspections/capitalization/MultipleReturns.java

@@ -0,0 +1,9 @@
+import org.jetbrains.annotations.Nls;
+
+class MultipleReturns  {
+    @Nls(capitalization = Nls.Capitalization.Title)
+    public String getName() {
+        if (true) return <warning descr="String 'Foo bar' is not properly capitalized. It should have title capitalization">"Foo bar"</warning>;
+        else return <warning descr="String 'Foo boo' is not properly capitalized. It should have title capitalization">"Foo <caret>boo"</warning>;
+    }
+}

plugins/devkit/testData/inspections/capitalization/MultipleReturns_after.java

@@ -0,0 +1,9 @@
+import org.jetbrains.annotations.Nls;
+
+class MultipleReturns  {
+    @Nls(capitalization = Nls.Capitalization.Title)
+    public String getName() {
+        if (true) return "Foo bar";
+        else return "Foo Boo";
+    }
+}

plugins/devkit/testSources/inspections/CapitalizationInspectionTest.java

@@ -33,6 +33,10 @@ public class CapitalizationInspectionTest extends LightCodeInsightFixtureTestCas
     doTest();
   }
 
+  public void testMultipleReturns() throws Exception {
+    doTest();
+  }
+
   public void doTest() {
     myFixture.testHighlighting(getTestName(false) + ".java");
     final IntentionAction action = myFixture.filterAvailableIntentions("Properly capitalize").get(0);
@@ -51,7 +55,6 @@ public class CapitalizationInspectionTest extends LightCodeInsightFixtureTestCas
     myFixture.addClass("package com.intellij.codeInspection; public class CommonProblemDescriptor {}");
     myFixture.addClass("package com.intellij.codeInspection; public class QuickFix {}");
     myFixture.enableInspections(TitleCapitalizationInspection.class);
-
   }