unworking version

coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/HandleSafely.kt

@@ -1,9 +1,9 @@
 package dev.inmo.micro_utils.coroutines
 
-import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.*
-import kotlinx.coroutines.supervisorScope
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.coroutineContext
+import kotlin.reflect.KClass
 
 typealias ExceptionHandler<T> = suspend (Throwable) -> T
 
@@ -16,13 +16,24 @@ var defaultSafelyExceptionHandler: ExceptionHandler<Nothing> = { throw it }
  * Key for [SafelyExceptionHandler] which can be used in [CoroutineContext.get] to get current default
  * [SafelyExceptionHandler]
  */
-class SafelyExceptionHandlerKey<T> : CoroutineContext.Key<SafelyExceptionHandler<T>>
+class SafelyExceptionHandlerKey<T>() : CoroutineContext.Key<SafelyExceptionHandler<T>>
+private val nothingSafelyEceptionHandlerKey = SafelyExceptionHandlerKey<Nothing>()
+private val unitSafelyEceptionHandlerKey = SafelyExceptionHandlerKey<Unit>()
 
+private val exceptionHandlersKeysCache = mutableMapOf<>()
 /**
  * Shortcut for creating instance of [SafelyExceptionHandlerKey]
  */
 @Suppress("NOTHING_TO_INLINE")
-inline fun <T> safelyExceptionHandlerKey() = SafelyExceptionHandlerKey<T>()
+inline fun <T> safelyExceptionHandlerKey() = SafelyExceptionHandlerKey(T::class)
+
+/**
+ * Shortcut for getting instance of [SafelyExceptionHandler] from current [coroutineContext]
+ */
+@Suppress("NOTHING_TO_INLINE")
+suspend inline fun <reified T : Any> safelyExceptionHandler() = coroutineContext[safelyExceptionHandlerKey<T>()]
+@Suppress("NOTHING_TO_INLINE")
+inline fun <reified T : Any> ExceptionHandler<T>.safelyExceptionHandler() = SafelyExceptionHandler(this, T::class)
 
 /**
  * Wrapper for [ExceptionHandler] which can be used in [CoroutineContext] to set local (for [CoroutineContext]) default
@@ -31,11 +42,11 @@ inline fun <T> safelyExceptionHandlerKey() = SafelyExceptionHandlerKey<T>()
  * @see SafelyExceptionHandlerKey
  * @see ExceptionHandler
  */
-class SafelyExceptionHandler<T>(
+class SafelyExceptionHandler<T : Any>(
-    val handler: ExceptionHandler<T>
+    val handler: ExceptionHandler<T>,
+    returnKClass: KClass<T>
 ) : CoroutineContext.Element {
-
+    override val key: CoroutineContext.Key<*> = SafelyExceptionHandlerKey(returnKClass)
-    override val key: CoroutineContext.Key<*> = safelyExceptionHandlerKey<T>()
 }
 
 /**
@@ -56,18 +67,28 @@ class SafelyExceptionHandler<T>(
  * @see SafelyExceptionHandler
  */
 suspend inline fun <T> safely(
-    noinline onException: ExceptionHandler<T> = defaultSafelyExceptionHandler,
+    onException: ExceptionHandler<T> = defaultSafelyExceptionHandler,
-    noinline block: suspend CoroutineScope.() -> T
+    block: suspend CoroutineScope.() -> T
 ): T {
+    val contextHandler = if (onException === defaultSafelyExceptionHandler) {
+        coroutineContext[nothingSafelyEceptionHandlerKey] ?:
+        safelyExceptionHandler<Unit>() ?.let { unitHandler ->
+            val handler = unitHandler.handler
+            SafelyExceptionHandler<T> {
+                handler(it)
+                onException(it)
+            }
+        } ?:
+        onException.safelyExceptionHandler()
+    } else {
+        onException.safelyExceptionHandler()
+    }
     return try {
-        supervisorScope(block)
+        withContext(contextHandler) {
-    } catch (e: Throwable) {
+            supervisorScope(block)
-        val handler = if (onException == defaultSafelyExceptionHandler) {
-            coroutineContext[safelyExceptionHandlerKey<T>()] ?.handler ?: onException
-        } else {
-            onException
         }
-        handler(e)
+    } catch (e: Throwable) {
+        contextHandler.handler(e)
     }
 }
 

coroutines/src/jvmTest/kotlin/dev/inmo/micro_utils/coroutines/HandleSafelyCoroutineContextTest.kt

@@ -0,0 +1,32 @@
+package dev.inmo.micro_utils.coroutines
+
+import kotlinx.coroutines.*
+import kotlin.test.Test
+
+class HandleSafelyCoroutineContextTest {
+    @Test
+    fun testHandleSafelyCoroutineContext() {
+        val scope = CoroutineScope(Dispatchers.Default)
+        var contextHandlerHappen = false
+        var localHandlerHappen = false
+        val contextHandler: ExceptionHandler<Unit> = {
+            contextHandlerHappen = true
+        }
+        val checkJob = scope.launch(Dispatchers.Default + contextHandler.safelyExceptionHandler()) {
+            safely {
+                safely(
+                    {
+                        localHandlerHappen = true
+                    }
+                ) {
+                    error("That must happen :)")
+                }
+                println(coroutineContext)
+                error("That must happen too:)")
+            }
+        }
+        launchSynchronously { checkJob.join() }
+        assert(contextHandlerHappen)
+        assert(localHandlerHappen)
+    }
+}