Feat: rememberSaveableWrapper gives Python Compose's rememberSaveable

docs/SPEC.md

@@ -400,6 +400,15 @@ All in `PM/commonTest`, so they run wherever the interpreter loads.
   that root is replaced or the composition is disposed. The entry point names no library.
   `Status: implemented` on desktop, `ksp-fixtures/compose/.../PythonAppViewRenderTest.kt`; the module
   compiles for Android, nothing runs there yet.
+- **B-9** A Python root keeps state across the recreation of its composition with
+  `python.multiplatform.compose.rememberSaveableWrapper(initial)` (the pycomposeui function of that name),
+  callable from Python inside a composition. A Python `int` (32 bit) becomes `rememberSaveable { mutableIntStateOf }`,
+  a larger one `mutableLongStateOf`, a `float` `mutableDoubleStateOf`, a `bool` or `str` `mutableStateOf`; the
+  returned `MutableState` is read and written through `.value`. A value written from Python survives
+  `SaveableStateRegistry.performSave()` and a new composition over a registry built from the saved map; without
+  a registry the new composition starts at `initial`. Any other value (a Python `list`) raises an error that
+  names its type. `Status: implemented` on desktop (#174),
+  `ksp-fixtures/compose/.../RememberSaveableRenderTest.kt`.
 - **B-7** The plugin generates `.pyi` stubs for the Kotlin-named modules only, under Kotlin names
   plus their U-12 Pythonic aliases (`fill_max_width = fillMaxWidth` beside each module `def`, a second
   `ClassVar` or `@property` on a receiver's class, a second member of a callable module), parameters

ksp-fixtures/compose/build.gradle.kts

@@ -111,6 +111,10 @@ pythonBindings {
             //    `TextFieldStateRenderTest` is what each one is for.
             "androidx.compose.foundation.text.input.TextFieldState",
             "androidx.compose.foundation.text.input.TextFieldStateKt",
+            // 8. `rememberSaveableWrapper` (issue #174): the one function of `python-multiplatform-compose`
+            //    Python calls inside a composition. Its file facade, not the package, so that the host
+            //    composables (`PythonAppView` and friends) stay out of the table.
+            "python.multiplatform.compose.RememberSaveableKt",
         ),
     )
     // Stubs on: `generatePythonStubs` over the real Compose jars is what `tools/stubs/check-stubs.sh` (the mypy

ksp-fixtures/compose/src/desktopTest/kotlin/fixture/compose/RememberSaveableRenderTest.kt

@@ -0,0 +1,185 @@
+package fixture.compose
+
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.saveable.LocalSaveableStateRegistry
+import androidx.compose.runtime.saveable.SaveableStateRegistry
+import androidx.compose.runtime.snapshots.Snapshot
+import androidx.compose.ui.ImageComposeScene
+import androidx.compose.ui.unit.Density
+import python.multiplatform.compose.PythonWidget
+import python.multiplatform.ffi.Python3
+import python.multiplatform.ffi.pythonx.PythonxAdapter
+import python.multiplatform.ffi.upcall.UpcallBootstrap
+import python.multiplatform.generated.artifacts.ArtifactTable
+import python.multiplatform.reflection.UpcallTable
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * **`rememberSaveableWrapper` from Python** (issue #174): the original pycomposeui's
+ * `rememberSaveableWrapper(init, type)`, as `python.multiplatform.compose.rememberSaveableWrapper(initial)`,
+ * reached through the walk (`RememberSaveableKt`) and called inside a Python root.
+ *
+ * The root asks for one state per supported type (Int, Long, Double, Boolean, String), draws them, and
+ * keeps them in `_rs_states` so the test can write `.value` from Python. The scene is hosted inside
+ * `CompositionLocalProvider(LocalSaveableStateRegistry provides registry)`; the test writes, saves the
+ * registry, closes the scene, builds a new registry from the saved map and a new scene over the same
+ * root, and reads the values the new states hold.
+ *
+ * ### Red before the implementation
+ *
+ * `from python.multiplatform.compose import rememberSaveableWrapper` is an `ImportError` raised by the
+ * Python root when the first frame composes, so each test fails there and not at an assertion on a
+ * value, which is what tells that red from a regression.
+ */
+class RememberSaveableRenderTest {
+
+    @BeforeTest
+    fun installProducers() {
+        Python3.initialize(silent = true)
+        UpcallTable.clear()
+        UpcallTable.install(ArtifactTable.fragments)
+        check(UpcallBootstrap.publishToGlobals()) { "UpcallBootstrap.publishToGlobals() failed" }
+        PythonxAdapter.install(PYTHONX_RAW_VALUE_CLASSES)
+        Python3.exec(ROOTS)
+    }
+
+    @AfterTest
+    fun cleanup() {
+        Python3.exec(
+            """
+            for _name in [n for n in globals() if n.startswith('_rs_')]:
+                del globals()[_name]
+            """.trimIndent(),
+        )
+        UpcallTable.clear()
+    }
+
+    /** Each supported type: initial value, written value, saved, restored into a new scene. */
+    @Test
+    fun everySupportedTypeSurvivesASaveAndARestoreThroughTheRegistry() {
+        val saved = firstScene(canBeSaved = true)
+        println("rememberSaveableWrapper: saved ${saved.size} entries")
+        assertTrue(saved.isNotEmpty(), "nothing was saved, so the states are not registered with the registry")
+
+        runScene(SaveableStateRegistry(restoredValues = saved, canBeSaved = { true }))
+        assertEquals(WRITTEN, snapshotOfStates(), "the restored states do not hold the values saved")
+        assertEquals(
+            listOf("int", "int", "float", "bool", "str"),
+            pyStr("[type(_rs_states[k].value).__name__ for k in 'ilfbs']").let { parseList(it) },
+        )
+    }
+
+    /**
+     * The check that can fail: the same flow hosted without a registry (so nothing can be saved) shows
+     * the initial values again.
+     */
+    @Test
+    fun withoutASaveableRegistryTheRecreatedSceneShowsTheInitialValues() {
+        val saved = firstScene(canBeSaved = null)
+        assertTrue(saved.isEmpty(), "a scene without a registry still saved $saved")
+
+        runScene(null)
+        assertEquals(INITIAL, snapshotOfStates(), "values came back although nothing was saved")
+    }
+
+    /** A value of any other type is an error that names the type. */
+    @Test
+    fun anUnsupportedTypeRaisesAnErrorNamingTheType() {
+        val scene = ImageComposeScene(width = 200, height = 60, density = Density(1f)) {
+            CompositionLocalProvider(LocalSaveableStateRegistry provides SaveableStateRegistry(null) { true }) {
+                PythonWidget("_rs_root_bad", moduleName = "__main__")
+            }
+        }
+        try {
+            scene.render()
+        } finally {
+            scene.close()
+        }
+        val message = pyStr("_rs_error[0]")
+        println("rememberSaveableWrapper unsupported: $message")
+        assertTrue("list" in message, "the error does not name the Python list type: $message")
+    }
+
+    private fun firstScene(canBeSaved: Boolean?): Map<String, List<Any?>> {
+        val registry = canBeSaved?.let { SaveableStateRegistry(restoredValues = null, canBeSaved = { _ -> it }) }
+        val scene = host(registry)
+        try {
+            scene.render()
+            assertEquals(INITIAL, snapshotOfStates(), "the states did not start at the values Python passed")
+            Python3.exec("_rs_states['i'].value = 8; _rs_states['l'].value = 6000000000; _rs_states['f'].value = 2.5; _rs_states['b'].value = False; _rs_states['s'].value = 'xyz'")
+            Snapshot.sendApplyNotifications()
+            scene.render()
+            assertEquals(WRITTEN, snapshotOfStates(), "the write from Python was not read back")
+            return registry?.performSave() ?: emptyMap()
+        } finally {
+            scene.close()
+        }
+    }
+
+    private fun runScene(registry: SaveableStateRegistry?) {
+        Python3.exec("_rs_states.clear()")
+        val scene = host(registry)
+        try {
+            scene.render()
+        } finally {
+            scene.close()
+        }
+    }
+
+    private fun host(registry: SaveableStateRegistry?) =
+        ImageComposeScene(width = 300, height = 200, density = Density(1f)) {
+            if (registry == null) {
+                PythonWidget("_rs_root", moduleName = "__main__")
+            } else {
+                CompositionLocalProvider(LocalSaveableStateRegistry provides registry) {
+                    PythonWidget("_rs_root", moduleName = "__main__")
+                }
+            }
+        }
+
+    private fun snapshotOfStates(): String = pyStr("[_rs_states[k].value for k in 'ilfbs']")
+
+    private fun parseList(text: String): List<String> =
+        text.removePrefix("[").removeSuffix("]").split(", ").map { it.trim('\'') }
+
+    private fun pyStr(expression: String): String = Python3.import("__main__").getAttr("__dict__").let { globals ->
+        Python3.eval(expression, PY_EVAL_INPUT, globals, globals).toString()
+    }
+
+    private companion object {
+        /** `Py_eval_input`. */
+        const val PY_EVAL_INPUT = 258
+        const val INITIAL = "[7, 5000000000, 1.5, True, 'abc']"
+        const val WRITTEN = "[8, 6000000000, 2.5, False, 'xyz']"
+
+        val ROOTS = """
+            _rs_states = {}
+            _rs_error = []
+
+            def _rs_root():
+                from python.multiplatform.compose import rememberSaveableWrapper
+                from androidx.compose.foundation.layout import Column
+                from androidx.compose.material3 import Text
+                _rs_states['i'] = rememberSaveableWrapper(7)
+                _rs_states['l'] = rememberSaveableWrapper(5000000000)
+                _rs_states['f'] = rememberSaveableWrapper(1.5)
+                _rs_states['b'] = rememberSaveableWrapper(True)
+                _rs_states['s'] = rememberSaveableWrapper('abc')
+                def _rs_body():
+                    for _k in 'ilfbs':
+                        Text(str(_rs_states[_k].value))
+                Column(content=_rs_body)
+
+            def _rs_root_bad():
+                from python.multiplatform.compose import rememberSaveableWrapper
+                try:
+                    rememberSaveableWrapper([1, 2])
+                except BaseException as _e:
+                    _rs_error.append(type(_e).__name__ + ': ' + str(_e))
+        """.trimIndent()
+    }
+}

python-multiplatform-compose/build.gradle.kts

@@ -36,6 +36,7 @@ kotlin {
             api(projects.pythonMultiplatform)
             api(compose.runtime)
             api(compose.ui)
+            api(compose.runtimeSaveable)
         }
     }
 }

python-multiplatform-compose/src/commonMain/kotlin/python/multiplatform/compose/RememberSaveable.kt

@@ -0,0 +1,53 @@
+package python.multiplatform.compose
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.MutableState
+import androidx.compose.runtime.mutableDoubleStateOf
+import androidx.compose.runtime.mutableIntStateOf
+import androidx.compose.runtime.mutableLongStateOf
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.saveable.rememberSaveable
+import python.multiplatform.ffi.PyObject
+
+/**
+ * `rememberSaveable { mutableStateOf(initial) }` for a value Python holds, so that a Python screen's
+ * state survives what `rememberSaveable` survives (recreating an activity, process death with a
+ * saved instance state, a `SaveableStateRegistry` saved and restored).
+ *
+ * Follows `rememberSaveableWrapper(init, type)` of the original pycomposeui
+ * (`Runtime.android.kt`), which dispatched on a type name Python passed along. Here the value alone
+ * is passed: a Python scalar written into a `kotlin.Any?` slot arrives boxed, so its runtime type is
+ * the type name.
+ *
+ * | Python value | Kotlin type | State created |
+ * |---|---|---|
+ * | `int` that fits 32 bits | `Int` | `mutableIntStateOf` |
+ * | larger `int` | `Long` | `mutableLongStateOf` |
+ * | `float` | `Double` | `mutableDoubleStateOf` |
+ * | `bool` | `Boolean` | `mutableStateOf` |
+ * | `str` | `String` | `mutableStateOf` |
+ *
+ * Any other value throws [IllegalArgumentException] naming its type (`list` for a Python list): like the original, only these
+ * are supported, because only these are saveable without a custom `Saver`.
+ *
+ * The returned state is Compose's own: Python reads and writes its `value`, and a write is saved by
+ * the next `SaveableStateRegistry.performSave`. Call it inside a composition only.
+ */
+@Composable
+fun rememberSaveableWrapper(initial: Any?): MutableState<*> = when (initial) {
+    is Int -> rememberSaveable { mutableIntStateOf(initial) }
+    is Long -> rememberSaveable { mutableLongStateOf(initial) }
+    is Double -> rememberSaveable { mutableDoubleStateOf(initial) }
+    is Boolean -> rememberSaveable { mutableStateOf(initial) }
+    is String -> rememberSaveable { mutableStateOf(initial) }
+    else -> throw IllegalArgumentException(
+        "rememberSaveableWrapper supports Int, Long, Double, Boolean and String, not " + typeNameOf(initial),
+    )
+}
+
+/** The type of [value] as Python spells it for a Python object (`list`), as Kotlin does otherwise. */
+private fun typeNameOf(value: Any?): String = when (value) {
+    null -> "null"
+    is PyObject -> value.Type.getAttr("__name__").toString()
+    else -> value::class.qualifiedName ?: value::class.simpleName ?: "unknown"
+}