Added //sampleStart //sampleEnd to samples on kotlin website

4 files changed, 30 insertions, 13 deletions

diff --git a/core/src/main/kotlin/Formats/KotlinWebsiteFormatService.kt b/core/src/main/kotlin/Formats/KotlinWebsiteFormatService.kt
index edae9976..ac70bb02 100644
--- a/core/src/main/kotlin/Formats/KotlinWebsiteFormatService.kt
+++ b/core/src/main/kotlin/Formats/KotlinWebsiteFormatService.kt
@@ -53,9 +53,14 @@ class KotlinWebsiteOutputBuilder(to: StringBuilder,
         }
     }
 
-    override fun appendSampleBlockCode(language: String, body: () -> Unit) {
+    override fun appendSampleBlockCode(language: String, imports: () -> Unit, body: () -> Unit) {
         div(to, "sample", true) {
-            appendBlockCode(language, body)
+            appendBlockCode(language) {
+                imports()
+                wrap("\nfun main(args: Array<String>) {", "}") {
+                    wrap("\n//sampleStart\n", "\n//sampleEnd\n", body)
+                }
+            }
         }
     }
 
diff --git a/core/src/main/kotlin/Formats/StructuredFormatService.kt b/core/src/main/kotlin/Formats/StructuredFormatService.kt
index 3bd06130..183282ea 100644
--- a/core/src/main/kotlin/Formats/StructuredFormatService.kt
+++ b/core/src/main/kotlin/Formats/StructuredFormatService.kt
@@ -47,7 +47,7 @@ abstract class StructuredOutputBuilder(val to: StringBuilder,
 
     protected abstract fun ensureParagraph()
 
-    open fun appendSampleBlockCode(language: String, body: () -> Unit) = appendBlockCode(language, body)
+    open fun appendSampleBlockCode(language: String, imports: () -> Unit, body: () -> Unit) = appendBlockCode(language, body)
     abstract fun appendBlockCode(language: String, body: () -> Unit)
     abstract fun appendHeader(level: Int = 1, body: () -> Unit)
     abstract fun appendParagraph(body: () -> Unit)
@@ -146,19 +146,19 @@ abstract class StructuredOutputBuilder(val to: StringBuilder,
 
             is ContentBlockSampleCode, is ContentBlockCode -> {
                 content as ContentBlockCode
-                fun appendBlockCodeContent() {
-                    for ((index, contentNode) in content.children.withIndex()) {
+                fun ContentBlockCode.appendBlockCodeContent() {
+                    for ((index, contentNode) in this.children.withIndex()) {
                         appendContent(contentNode)
-                        if (index < content.children.size - 1) {
+                        if (index < this.children.size - 1) {
                             to.append("\n")
                         }
                     }
                 }
                 when (content) {
                     is ContentBlockSampleCode ->
-                        appendSampleBlockCode(content.language, ::appendBlockCodeContent)
+                        appendSampleBlockCode(content.language, content.importsBlock::appendBlockCodeContent, { content.appendBlockCodeContent() })
                     is ContentBlockCode ->
-                        appendBlockCode(content.language, ::appendBlockCodeContent)
+                        appendBlockCode(content.language, { content.appendBlockCodeContent() })
                 }
             }
             is ContentHeading -> appendHeader(content.level) { appendContent(content.children) }
diff --git a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
index 2fc268e8..532b186c 100644
--- a/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
+++ b/core/src/main/kotlin/Kotlin/DescriptorDocumentationParser.kt
@@ -18,6 +18,7 @@ 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.DescriptorUtils
 import org.jetbrains.kotlin.resolve.annotations.argumentValue
@@ -148,7 +149,7 @@ 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().let { it.append(ContentText("//Missing function name in @sample")); it }
+            return ContentBlockSampleCode().apply { append(ContentText("//Missing function name in @sample")) }
         }
         val scope = getKDocLinkResolutionScope(resolutionFacade, descriptor)
         val rootPackage = resolutionFacade.moduleDescriptor.getPackage(FqName.ROOT)
@@ -156,12 +157,12 @@ class DescriptorDocumentationParser
         val symbol = resolveInScope(functionName, scope) ?: resolveInScope(functionName, rootScope)
         if (symbol == null) {
             logger.warn("Unresolved function $functionName in @sample in ${descriptor.signature()}")
-            return ContentBlockSampleCode().let { it.append(ContentText("//Unresolved: $functionName")); it }
+            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().let { it.append(ContentText("//Source not found: $functionName")); it }
+            return ContentBlockSampleCode().apply { append(ContentText("//Source not found: $functionName")) }
         }
 
         val text = when (psiElement) {
@@ -178,7 +179,18 @@ class DescriptorDocumentationParser
         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().let { it.append(ContentText(finalText)); it }
+
+        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? {
diff --git a/core/src/main/kotlin/Model/Content.kt b/core/src/main/kotlin/Model/Content.kt
index fcf94c12..ecf75b20 100644
--- a/core/src/main/kotlin/Model/Content.kt
+++ b/core/src/main/kotlin/Model/Content.kt
@@ -83,7 +83,7 @@ class ContentStrong() : ContentBlock()
 class ContentStrikethrough() : ContentBlock()
 class ContentCode() : ContentBlock()
 open class ContentBlockCode(val language: String = "") : ContentBlock()
-class ContentBlockSampleCode(language: String = "kotlin") : ContentBlockCode(language)
+class ContentBlockSampleCode(language: String = "kotlin", val importsBlock: ContentBlockCode = ContentBlockCode(language)) : ContentBlockCode(language)
 
 abstract class ContentNodeLink() : ContentBlock() {
     abstract val node: DocumentationNode?