Replacing assertPrints, assertTrue to println for kotlin-website samples, Added SampleProcessingService to FormatDescriptor

10 files changed, 221 insertions, 81 deletions

diff --git a/core/src/main/kotlin/Formats/FormatDescriptor.kt b/core/src/main/kotlin/Formats/FormatDescriptor.kt
index e384f223..58c9472b 100644
--- a/core/src/main/kotlin/Formats/FormatDescriptor.kt
+++ b/core/src/main/kotlin/Formats/FormatDescriptor.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.dokka.Formats
 
+import Samples.SampleProcessingService
 import org.jetbrains.dokka.*
 import kotlin.reflect.KClass
 
@@ -9,4 +10,5 @@ interface FormatDescriptor {
     val generatorServiceClass: KClass<out Generator>
     val packageDocumentationBuilderClass: KClass<out PackageDocumentationBuilder>
     val javaDocumentationBuilderClass: KClass<out JavaDocumentationBuilder>
+    val sampleProcessingService: KClass<out SampleProcessingService>
 }
diff --git a/core/src/main/kotlin/Formats/StandardFormats.kt b/core/src/main/kotlin/Formats/StandardFormats.kt
index fdc8eb9e..45a653b0 100644
--- a/core/src/main/kotlin/Formats/StandardFormats.kt
+++ b/core/src/main/kotlin/Formats/StandardFormats.kt
@@ -1,5 +1,8 @@
 package org.jetbrains.dokka.Formats
 
+import Samples.DefaultSampleProcessingService
+import Samples.KotlinWebsiteSampleProcessingService
+import Samples.SampleProcessingService
 import org.jetbrains.dokka.*
 import kotlin.reflect.KClass
 
@@ -9,6 +12,7 @@ abstract class KotlinFormatDescriptorBase : FormatDescriptor {
 
     override val generatorServiceClass = FileGenerator::class
     override val outlineServiceClass: KClass<out OutlineFormatService>? = null
+    override val sampleProcessingService: KClass<out SampleProcessingService> = DefaultSampleProcessingService::class
 }
 
 class HtmlFormatDescriptor : KotlinFormatDescriptorBase() {
@@ -22,11 +26,13 @@ class HtmlAsJavaFormatDescriptor : FormatDescriptor {
     override val generatorServiceClass = FileGenerator::class
     override val packageDocumentationBuilderClass = KotlinAsJavaDocumentationBuilder::class
     override val javaDocumentationBuilderClass = JavaPsiDocumentationBuilder::class
+    override val sampleProcessingService: KClass<out SampleProcessingService> = DefaultSampleProcessingService::class
 }
 
 class KotlinWebsiteFormatDescriptor : KotlinFormatDescriptorBase() {
     override val formatServiceClass = KotlinWebsiteFormatService::class
     override val outlineServiceClass = YamlOutlineService::class
+    override val sampleProcessingService: KClass<out SampleProcessingService> = KotlinWebsiteSampleProcessingService::class
 }
 
 class JekyllFormatDescriptor : KotlinFormatDescriptorBase() {
diff --git a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
index 532b186c..a7aa6452 100644
--- a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
+++ b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.dokka.Kotlin
 
+import Samples.SampleProcessingService
 import com.google.inject.Inject
 import com.intellij.psi.PsiDocCommentOwner
 import com.intellij.psi.PsiNamedElement
@@ -17,6 +18,7 @@ import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
@@ -33,7 +35,8 @@ class DescriptorDocumentationParser
                              val logger: DokkaLogger,
                              val linkResolver: DeclarationLinkResolver,
                              val resolutionFacade: DokkaResolutionFacade,
-                             val refGraph: NodeReferenceGraph)
+                             val refGraph: NodeReferenceGraph,
+                             val sampleService: SampleProcessingService)
 {
     fun parseDocumentation(descriptor: DeclarationDescriptor, inline: Boolean = false): Content =
             parseDocumentationAndDetails(descriptor, inline).first
@@ -64,7 +67,7 @@ class DescriptorDocumentationParser
             tags.forEach {
                 when (it.knownTag) {
                     KDocKnownTag.SAMPLE ->
-                        content.append(functionBody(descriptor, it.getSubjectName()))
+                        content.append(sampleService.resolveSample(descriptor, it.getSubjectName()))
                     KDocKnownTag.SEE ->
                         content.addTagToSeeAlso(descriptor, it)
                     else -> {
@@ -146,80 +149,4 @@ class DescriptorDocumentationParser
         }
     }
 
-    private fun functionBody(descriptor: DeclarationDescriptor, functionName: String?): ContentNode {
-        if (functionName == null) {
-            logger.warn("Missing function name in @sample in ${descriptor.signature()}")
-            return ContentBlockSampleCode().apply { append(ContentText("//Missing function name in @sample")) }
-        }
-        val scope = getKDocLinkResolutionScope(resolutionFacade, descriptor)
-        val rootPackage = resolutionFacade.moduleDescriptor.getPackage(FqName.ROOT)
-        val rootScope = rootPackage.memberScope
-        val symbol = resolveInScope(functionName, scope) ?: resolveInScope(functionName, rootScope)
-        if (symbol == null) {
-            logger.warn("Unresolved function $functionName in @sample in ${descriptor.signature()}")
-            return ContentBlockSampleCode().apply { append(ContentText("//Unresolved: $functionName")) }
-        }
-        val psiElement = DescriptorToSourceUtils.descriptorToDeclaration(symbol)
-        if (psiElement == null) {
-            logger.warn("Can't find source for function $functionName in @sample in ${descriptor.signature()}")
-            return ContentBlockSampleCode().apply { append(ContentText("//Source not found: $functionName")) }
-        }
-
-        val text = when (psiElement) {
-            is KtDeclarationWithBody -> ContentBlockCode().let {
-                val bodyExpression = psiElement.bodyExpression
-                when (bodyExpression) {
-                    is KtBlockExpression -> bodyExpression.text.removeSurrounding("{", "}")
-                    else -> bodyExpression!!.text
-                }
-            }
-            else -> psiElement.text
-        }
-
-        val lines = text.trimEnd().split("\n".toRegex()).toTypedArray().filterNot(String::isEmpty)
-        val indent = lines.map { it.takeWhile(Char::isWhitespace).count() }.min() ?: 0
-        val finalText = lines.map { it.drop(indent) }.joinToString("\n")
-
-        val psiFile = psiElement.containingFile
-        val importsBlock = if (psiFile is KtFile) {
-            ContentBlockCode("kotlin").apply {
-                append(ContentText(psiFile.importList?.text ?: ""))
-            }
-        } else {
-            ContentBlockCode("")
-        }
-
-
-        return ContentBlockSampleCode(importsBlock = importsBlock).apply { append(ContentText(finalText)) }
-    }
-
-    private fun resolveInScope(functionName: String, scope: ResolutionScope): DeclarationDescriptor? {
-        var currentScope = scope
-        val parts = functionName.split('.')
-
-        var symbol: DeclarationDescriptor? = null
-
-        for (part in parts) {
-            // short name
-            val symbolName = Name.identifier(part)
-            val partSymbol = currentScope.getContributedDescriptors(DescriptorKindFilter.ALL, { it == symbolName })
-                    .filter { it.name == symbolName }
-                    .firstOrNull()
-
-            if (partSymbol == null) {
-                symbol = null
-                break
-            }
-            @Suppress("IfThenToElvis")
-            currentScope = if (partSymbol is ClassDescriptor)
-                partSymbol.defaultType.memberScope
-            else if (partSymbol is PackageViewDescriptor)
-                partSymbol.memberScope
-            else
-                getKDocLinkResolutionScope(resolutionFacade, partSymbol)
-            symbol = partSymbol
-        }
-
-        return symbol
-    }
 }
diff --git a/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt b/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt
new file mode 100644
index 00000000..26841ea6
--- /dev/null
+++ b/core/src/main/kotlin/Samples/DefaultSampleProcessingService.kt
@@ -0,0 +1,106 @@
+package Samples
+
+import com.google.inject.Inject
+import com.intellij.psi.PsiElement
+import org.jetbrains.dokka.*
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.PackageViewDescriptor
+import org.jetbrains.kotlin.idea.kdoc.getKDocLinkResolutionScope
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDeclarationWithBody
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.resolve.scopes.ResolutionScope
+
+
+open class DefaultSampleProcessingService
+@Inject constructor(val options: DocumentationOptions,
+                    val logger: DokkaLogger,
+                    val resolutionFacade: DokkaResolutionFacade)
+    : SampleProcessingService {
+
+    override fun resolveSample(descriptor: DeclarationDescriptor, functionName: String?): ContentNode {
+        if (functionName == null) {
+            logger.warn("Missing function name in @sample in ${descriptor.signature()}")
+            return ContentBlockSampleCode().apply { append(ContentText("//Missing function name in @sample")) }
+        }
+        val scope = getKDocLinkResolutionScope(resolutionFacade, descriptor)
+        val rootPackage = resolutionFacade.moduleDescriptor.getPackage(FqName.ROOT)
+        val rootScope = rootPackage.memberScope
+        val symbol = resolveInScope(functionName, scope) ?: resolveInScope(functionName, rootScope)
+        if (symbol == null) {
+            logger.warn("Unresolved function $functionName in @sample in ${descriptor.signature()}")
+            return ContentBlockSampleCode().apply { append(ContentText("//Unresolved: $functionName")) }
+        }
+        val psiElement = DescriptorToSourceUtils.descriptorToDeclaration(symbol)
+        if (psiElement == null) {
+            logger.warn("Can't find source for function $functionName in @sample in ${descriptor.signature()}")
+            return ContentBlockSampleCode().apply { append(ContentText("//Source not found: $functionName")) }
+        }
+
+        val text = processSampleBody(psiElement)
+
+        val lines = text.trimEnd().split("\n".toRegex()).toTypedArray().filterNot(String::isEmpty)
+        val indent = lines.map { it.takeWhile(Char::isWhitespace).count() }.min() ?: 0
+        val finalText = lines.map { it.drop(indent) }.joinToString("\n")
+
+        return ContentBlockSampleCode(importsBlock = processImports(psiElement)).apply { append(ContentText(finalText)) }
+    }
+
+    protected open fun processSampleBody(psiElement: PsiElement): String = when (psiElement) {
+        is KtDeclarationWithBody -> {
+            val bodyExpression = psiElement.bodyExpression
+            when (bodyExpression) {
+                is KtBlockExpression -> bodyExpression.text.removeSurrounding("{", "}")
+                else -> bodyExpression!!.text
+            }
+        }
+        else -> psiElement.text
+    }
+
+    protected open fun processImports(psiElement: PsiElement): ContentBlockCode {
+        val psiFile = psiElement.containingFile
+        if (psiFile is KtFile) {
+            return ContentBlockCode("kotlin").apply {
+                append(ContentText(psiFile.importList?.text ?: ""))
+            }
+        } else {
+            return ContentBlockCode("")
+        }
+    }
+
+    private fun resolveInScope(functionName: String, scope: ResolutionScope): DeclarationDescriptor? {
+        var currentScope = scope
+        val parts = functionName.split('.')
+
+        var symbol: DeclarationDescriptor? = null
+
+        for (part in parts) {
+            // short name
+            val symbolName = Name.identifier(part)
+            val partSymbol = currentScope.getContributedDescriptors(DescriptorKindFilter.ALL, { it == symbolName })
+                    .filter { it.name == symbolName }
+                    .firstOrNull()
+
+            if (partSymbol == null) {
+                symbol = null
+                break
+            }
+            @Suppress("IfThenToElvis")
+            currentScope = if (partSymbol is ClassDescriptor)
+                partSymbol.defaultType.memberScope
+            else if (partSymbol is PackageViewDescriptor)
+                partSymbol.memberScope
+            else
+                getKDocLinkResolutionScope(resolutionFacade, partSymbol)
+            symbol = partSymbol
+        }
+
+        return symbol
+    }
+}
+
diff --git a/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
new file mode 100644
index 00000000..864734ea
--- /dev/null
+++ b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
@@ -0,0 +1,81 @@
+package Samples
+
+import com.google.inject.Inject
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import org.jetbrains.dokka.DocumentationOptions
+import org.jetbrains.dokka.DokkaLogger
+import org.jetbrains.dokka.DokkaResolutionFacade
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import java.util.*
+
+open class KotlinWebsiteSampleProcessingService
+@Inject constructor(options: DocumentationOptions,
+                    logger: DokkaLogger,
+                    resolutionFacade: DokkaResolutionFacade)
+    : DefaultSampleProcessingService(options, logger, resolutionFacade) {
+
+
+    fun buildReplacementData(psiElement: PsiElement): Map<TextRange, String> {
+
+        val result = TreeMap<TextRange, String>({ o1, o2 -> o1.startOffset.compareTo(o2.startOffset) })
+        fun convertAssertPrints(expression: KtCallExpression) {
+            val (argument, commentArgument) = expression.valueArguments
+            val comment = commentArgument.getArgumentExpression() as KtStringTemplateExpression
+            val commentText = comment.entries.joinToString("") { it.text }
+            result[expression.textRange] = "println(${argument.text}) //$commentText"
+        }
+
+        fun convertAssertTrue(expression: KtCallExpression) {
+            val (argument) = expression.valueArguments
+            result[expression.textRange] = "println(\"${argument.text} is \${${argument.text}}\") //true "
+        }
+
+        if (psiElement is KtElement) {
+            val visitor = object : KtTreeVisitor<Any>() {
+                override fun visitCallExpression(expression: KtCallExpression, data: Any?): Void? {
+                    when (expression.calleeExpression?.text) {
+                        "assertPrints" -> convertAssertPrints(expression)
+                        "assertTrue" -> convertAssertTrue(expression)
+                        else -> super.visitCallExpression(expression, data)
+                    }
+                    return null
+                }
+            }
+            psiElement.acceptChildren(visitor)
+        }
+        return result
+    }
+
+    private fun String.applyReplacements(baseOffset: Int, replacementData: Map<TextRange, String>): String {
+        val partsList = arrayListOf<String>()
+        var prevRange = TextRange(0, baseOffset)
+        for ((range, replacement) in replacementData) {
+            partsList.add(substring(prevRange.endOffset - baseOffset, range.startOffset - baseOffset))
+            partsList.add(replacement)
+            prevRange = range
+        }
+        partsList.add(substring(prevRange.endOffset - baseOffset))
+        return partsList.joinToString(separator = "")
+    }
+
+
+    override fun processSampleBody(psiElement: PsiElement): String {
+
+        val replacementData = buildReplacementData(psiElement)
+
+        return when (psiElement) {
+            is KtDeclarationWithBody -> {
+                val bodyExpression = psiElement.bodyExpression
+                val bodyExpressionText = bodyExpression!!.text.applyReplacements(bodyExpression.startOffset, replacementData)
+                when (bodyExpression) {
+                    is KtBlockExpression -> bodyExpressionText.removeSurrounding("{", "}")
+                    else -> bodyExpressionText
+                }
+            }
+            else -> psiElement.text.applyReplacements(psiElement.startOffset, replacementData)
+        }
+    }
+}
+
diff --git a/core/src/main/kotlin/Samples/SampleProcessingService.kt b/core/src/main/kotlin/Samples/SampleProcessingService.kt
new file mode 100644
index 00000000..61abda6d
--- /dev/null
+++ b/core/src/main/kotlin/Samples/SampleProcessingService.kt
@@ -0,0 +1,8 @@
+package Samples
+
+import org.jetbrains.dokka.ContentNode
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+
+interface SampleProcessingService {
+    fun resolveSample(descriptor: DeclarationDescriptor, functionName: String?): ContentNode
+}
\ No newline at end of file
diff --git a/core/src/main/kotlin/Utilities/DokkaModule.kt b/core/src/main/kotlin/Utilities/DokkaModule.kt
index 1eb82313..4b2a99d3 100644
--- a/core/src/main/kotlin/Utilities/DokkaModule.kt
+++ b/core/src/main/kotlin/Utilities/DokkaModule.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.dokka.Utilities
 
+import Samples.SampleProcessingService
 import com.google.inject.Binder
 import com.google.inject.Module
 import com.google.inject.Provider
@@ -44,6 +45,7 @@ class DokkaModule(val environment: AnalysisEnvironment,
         }
         binder.bind<PackageDocumentationBuilder>().to(descriptor.packageDocumentationBuilderClass.java)
         binder.bind<JavaDocumentationBuilder>().to(descriptor.javaDocumentationBuilderClass.java)
+        binder.bind<SampleProcessingService>().to(descriptor.sampleProcessingService.java)
 
         binder.bind<Generator>().to(descriptor.generatorServiceClass.java)
 
diff --git a/core/src/main/kotlin/javadoc/dokka-adapters.kt b/core/src/main/kotlin/javadoc/dokka-adapters.kt
index 23ee1702..e3a58331 100644
--- a/core/src/main/kotlin/javadoc/dokka-adapters.kt
+++ b/core/src/main/kotlin/javadoc/dokka-adapters.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.dokka.javadoc
 
+import Samples.DefaultSampleProcessingService
 import com.google.inject.Inject
 import com.sun.tools.doclets.formats.html.HtmlDoclet
 import org.jetbrains.dokka.*
@@ -27,4 +28,5 @@ class JavadocFormatDescriptor : FormatDescriptor {
     override val generatorServiceClass = JavadocGenerator::class
     override val packageDocumentationBuilderClass = KotlinAsJavaDocumentationBuilder::class
     override val javaDocumentationBuilderClass = JavaPsiDocumentationBuilder::class
+    override val sampleProcessingService = DefaultSampleProcessingService::class
 }
diff --git a/core/src/test/kotlin/TestAPI.kt b/core/src/test/kotlin/TestAPI.kt
index 508d28fc..7197d2c4 100644
--- a/core/src/test/kotlin/TestAPI.kt
+++ b/core/src/test/kotlin/TestAPI.kt
@@ -125,8 +125,9 @@ fun verifyOutput(roots: Array<ContentRoot>,
                  outputExtension: String,
                  withJdk: Boolean = false,
                  withKotlinRuntime: Boolean = false,
+                 format: String = "html",
                  outputGenerator: (DocumentationModule, StringBuilder) -> Unit) {
-    verifyModel(*roots, withJdk = withJdk, withKotlinRuntime = withKotlinRuntime) {
+    verifyModel(*roots, withJdk = withJdk, withKotlinRuntime = withKotlinRuntime, format = format) {
         verifyModelOutput(it, outputExtension, outputGenerator, roots.first().path)
     }
 }
@@ -147,8 +148,9 @@ fun verifyOutput(path: String,
                  outputExtension: String,
                  withJdk: Boolean = false,
                  withKotlinRuntime: Boolean = false,
+                 format: String = "html",
                  outputGenerator: (DocumentationModule, StringBuilder) -> Unit) {
-    verifyOutput(arrayOf(contentRootFromPath(path)), outputExtension, withJdk, withKotlinRuntime, outputGenerator)
+    verifyOutput(arrayOf(contentRootFromPath(path)), outputExtension, withJdk, withKotlinRuntime, format, outputGenerator)
 }
 
 fun verifyJavaOutput(path: String,
diff --git a/core/src/test/kotlin/format/KotlinWebSiteFormatTest.kt b/core/src/test/kotlin/format/KotlinWebSiteFormatTest.kt
index e7419ec8..e7677862 100644
--- a/core/src/test/kotlin/format/KotlinWebSiteFormatTest.kt
+++ b/core/src/test/kotlin/format/KotlinWebSiteFormatTest.kt
@@ -19,8 +19,12 @@ class KotlinWebSiteFormatTest {
         verifyKWSNodeByName("overloadGroup", "magic")
     }
 
+    @Test fun sampleWithAsserts() {
+        verifyKWSNodeByName("sampleWithAsserts", "a")
+    }
+
     private fun verifyKWSNodeByName(fileName: String, name: String) {
-        verifyOutput("testdata/format/website/$fileName.kt", ".md") { model, output ->
+        verifyOutput("testdata/format/website/$fileName.kt", ".md", format = "kotlin-website") { model, output ->
             kwsService.createOutputBuilder(output, tempLocation).appendNodes(model.members.single().members.filter { it.name == name })
         }
     }