Screenshotting functionality. (#110)

JetBrains · May 14, 2021 · 783ae27 · 783ae27
1 parent dfe5860
commit 783ae27
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 8 deletions.
diff --git a/samples/SkijaInjectSample/build.gradle.kts b/samples/SkijaInjectSample/build.gradle.kts
@@ -53,6 +53,7 @@ val additionalArguments = mutableMapOf<String, String>()
 val casualRun = tasks.named<JavaExec>("run") {
     systemProperty("skiko.fps.enabled", "true")
     systemProperty("skiko.hardwareInfo.enabled", "true")
+    jvmArgs?.add("-ea")
     // Use systemProperty("skiko.library.path", "/tmp") to test loader.
     System.getProperties().entries
         .associate {

diff --git a/samples/SkijaInjectSample/src/main/kotlin/SkijaInjectSample/App.kt b/samples/SkijaInjectSample/src/main/kotlin/SkijaInjectSample/App.kt
@@ -1,9 +1,7 @@
 package SkijaInjectSample
 
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.*
 import kotlinx.coroutines.swing.Swing
-import kotlinx.coroutines.yield
 import org.jetbrains.skija.*
 import org.jetbrains.skija.paragraph.FontCollection
 import org.jetbrains.skija.paragraph.ParagraphBuilder
@@ -12,12 +10,16 @@ import org.jetbrains.skija.paragraph.TextStyle
 import org.jetbrains.skiko.SkiaLayer
 import org.jetbrains.skiko.SkiaRenderer
 import org.jetbrains.skiko.SkiaWindow
+import org.jetbrains.skiko.toBufferedImage
 import java.awt.Dimension
 import java.awt.Toolkit
 import java.awt.event.*
 import javax.swing.*
 import kotlin.math.cos
 import kotlin.math.sin
+import java.io.File
+import java.nio.file.Files
+import javax.imageio.ImageIO
 
 fun main(args: Array<String>) {
     repeat(1) {
@@ -56,8 +58,25 @@ fun createWindow(title: String) = runBlocking(Dispatchers.Swing) {
         }
     })
 
+    val defaultScreenshotPath =
+        Files.createTempFile("compose_", ".png").toAbsolutePath().toString()
+    val miTakeScreenshot = JMenuItem("Take screenshot to $defaultScreenshotPath")
+    val ctrlS = KeyStroke.getKeyStroke(KeyEvent.VK_S, Toolkit.getDefaultToolkit().getMenuShortcutKeyMaskEx())
+    miTakeScreenshot.setAccelerator(ctrlS)
+    miTakeScreenshot.addActionListener(object : ActionListener {
+        override fun actionPerformed(actionEvent: ActionEvent?) {
+            val screenshot = window.layer.screenshot()!!
+            GlobalScope.launch(Dispatchers.IO) {
+                val image = screenshot.toBufferedImage()
+                ImageIO.write(image, "png", File(defaultScreenshotPath))
+                println("Saved to $defaultScreenshotPath")
+            }
+        }
+    })
+
     menu.add(miToggleFullscreen)
     menu.add(miFullscreenState)
+    menu.add(miTakeScreenshot)
 
     window.setJMenuBar(menuBar)
 

diff --git a/skiko/src/jvmMain/kotlin/org/jetbrains/skiko/Convertors.kt b/skiko/src/jvmMain/kotlin/org/jetbrains/skiko/Convertors.kt
@@ -0,0 +1,46 @@
+package org.jetbrains.skiko
+
+import org.jetbrains.skija.Bitmap
+import org.jetbrains.skija.ColorType
+import java.awt.Transparency
+import java.awt.color.ColorSpace
+import java.awt.image.BufferedImage
+import java.awt.image.ComponentColorModel
+import java.awt.image.DataBuffer
+import java.awt.image.Raster
+import java.nio.ByteBuffer
+
+private class DirectDataBuffer(val backing: ByteBuffer): DataBuffer(TYPE_BYTE, backing.limit()) {
+    override fun getElem(bank: Int, index: Int): Int {
+        return backing[index].toInt()
+    }
+    override fun setElem(bank: Int, index: Int, value: Int) {
+        throw UnsupportedOperationException("no write access")
+    }
+}
+
+fun Bitmap.toBufferedImage(): BufferedImage {
+    val pixels = this.peekPixels()
+    val order = when (this.colorInfo.colorType) {
+        ColorType.RGB_888X -> intArrayOf(0, 1, 2, 3)
+        ColorType.BGRA_8888 -> intArrayOf(2, 1, 0, 3)
+        else -> throw UnsupportedOperationException("unsupported color type ${this.colorInfo.colorType}")
+    }
+    val raster = Raster.createInterleavedRaster(
+        DirectDataBuffer(pixels!!),
+        this.width,
+        this.height,
+        this.width * 4,
+        4,
+        order,
+        null
+    )
+    val colorModel = ComponentColorModel(
+        ColorSpace.getInstance(ColorSpace.CS_sRGB),
+        true,
+        false,
+        Transparency.TRANSLUCENT,
+        DataBuffer.TYPE_BYTE
+    )
+   return BufferedImage(colorModel, raster!!, false, null)
+}
diff --git a/skiko/src/jvmMain/kotlin/org/jetbrains/skiko/SkiaLayer.kt b/skiko/src/jvmMain/kotlin/org/jetbrains/skiko/SkiaLayer.kt
@@ -4,11 +4,7 @@ import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.swing.Swing
 import kotlinx.coroutines.withContext
-import org.jetbrains.skija.Canvas
-import org.jetbrains.skija.ClipMode
-import org.jetbrains.skija.Picture
-import org.jetbrains.skija.PictureRecorder
-import org.jetbrains.skija.Rect
+import org.jetbrains.skija.*
 import org.jetbrains.skiko.context.ContextHandler
 import org.jetbrains.skiko.context.createContextHandler
 import org.jetbrains.skiko.redrawer.Redrawer
@@ -334,6 +330,28 @@ open class SkiaLayer(
         }
     }
 
+    // Captures current layer as bitmap.
+    fun screenshot(): Bitmap? {
+        return contextHandler?.let {
+            synchronized(pictureLock) {
+                val picture = picture
+                if (picture != null) {
+                    val store = Bitmap()
+                    val ci = ColorInfo(
+                        ColorType.BGRA_8888, ColorAlphaType.OPAQUE, ColorSpace.getSRGBLinear())
+                    store.imageInfo = ImageInfo(ci, picture.width, picture.height)
+                    store.allocN32Pixels(picture.width, picture.height)
+                    val canvas = Canvas(store)
+                    canvas.drawPicture(picture.instance)
+                    store.setImmutable()
+                    store
+                } else {
+                    null
+                }
+            }
+        }
+    }
+
     private fun Canvas.clipRectBy(rectangle: ClipRectangle) {
         val dpi = contentScale
         clipRect(