From b284bfc85fedaf9df8cd8f008788bb50e218917c Mon Sep 17 00:00:00 2001
From: Dmitry Avdeev <Dmitry.Avdeev@jetbrains.com>
Date: Tue, 23 Dec 2014 15:54:15 +0300
Subject: [PATCH] TitleCapitalizationInspection: multiple returns

---
 .../TitleCapitalizationInspection.java        | 19 +++++++++++--------
 .../capitalization/MultipleReturns.java       |  9 +++++++++
 .../capitalization/MultipleReturns_after.java |  9 +++++++++
 .../CapitalizationInspectionTest.java         |  5 ++++-
 4 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 plugins/devkit/testData/inspections/capitalization/MultipleReturns.java
 create mode 100644 plugins/devkit/testData/inspections/capitalization/MultipleReturns_after.java

diff --git a/plugins/devkit/src/inspections/TitleCapitalizationInspection.java b/plugins/devkit/src/inspections/TitleCapitalizationInspection.java
index 4b4061b15fcb..2de34c1856de 100644
--- a/plugins/devkit/src/inspections/TitleCapitalizationInspection.java
+++ b/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
diff --git a/plugins/devkit/testData/inspections/capitalization/MultipleReturns.java b/plugins/devkit/testData/inspections/capitalization/MultipleReturns.java
new file mode 100644
index 000000000000..2ba739b3bf20
--- /dev/null
+++ b/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>;
+    }
+}
\ No newline at end of file
diff --git a/plugins/devkit/testData/inspections/capitalization/MultipleReturns_after.java b/plugins/devkit/testData/inspections/capitalization/MultipleReturns_after.java
new file mode 100644
index 000000000000..eee08c8982db
--- /dev/null
+++ b/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";
+    }
+}
\ No newline at end of file
diff --git a/plugins/devkit/testSources/inspections/CapitalizationInspectionTest.java b/plugins/devkit/testSources/inspections/CapitalizationInspectionTest.java
index 8e4398c7511c..51b181a7ae84 100644
--- a/plugins/devkit/testSources/inspections/CapitalizationInspectionTest.java
+++ b/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);
-
   }