CommandLineToolkit

1
import Foundation
2

3
protocol ConsoleContextKey {
4
    associatedtype Value
5

6
    static var defaultValue: Value { get }
7
}
8

9
struct ConsoleContext {
10
    @TaskLocal
11
    static var current: ConsoleContext = .init()
12

13
    private var storage: [ObjectIdentifier: Any] = [:]
14

15
    subscript <Key: ConsoleContextKey>(_ key: Key.Type) -> Key.Value {
16
        get { storage[ObjectIdentifier(key)] as? Key.Value ?? key.defaultValue }
17
        set { storage[ObjectIdentifier(key)] = newValue }
18
    }
19
    
20
    func withUpdated<Value>(key: WritableKeyPath<ConsoleContext, Value>, value: Value) -> ConsoleContext {
21
        var newContext = self
22
        newContext[keyPath: key] = value
23
        return newContext
24
    }
25
}
26

27
extension TaskLocal where Value == ConsoleContext {
28
    func withUpdated<R, ContextValue>(
29
        key: WritableKeyPath<ConsoleContext, ContextValue>,
30
        value: ContextValue,
31
        operation: () throws -> R
32
    ) rethrows -> R {
33
        try self.withValue(wrappedValue.withUpdated(key: key, value: value), operation: operation)
34
    }
35
    
36
    func withUpdated<R, ContextValue>(
37
        key: WritableKeyPath<ConsoleContext, ContextValue>,
38
        value: ContextValue,
39
        operation: () async throws -> R
40
    ) async rethrows -> R {
41
        try await self.withValue(wrappedValue.withUpdated(key: key, value: value), operation: operation)
42
    }
43
}
44