From ca23e5d0ff02568f8bc5c66c3adea31517b06a56 Mon Sep 17 00:00:00 2001
From: Eugene Zolotarev <eugzol@gmail.com>
Date: Tue, 30 Jan 2024 16:42:57 +0300
Subject: [PATCH] Switch node fix

---
 lib/datacaster/switch_node.rb | 22 +++++++++++++++-------
 lib/datacaster/version.rb     |  2 +-
 spec/switch_node_spec.rb      | 10 ++++++++++
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/lib/datacaster/switch_node.rb b/lib/datacaster/switch_node.rb
index 1a107f9..65705b5 100644
--- a/lib/datacaster/switch_node.rb
+++ b/lib/datacaster/switch_node.rb
@@ -1,11 +1,15 @@
 module Datacaster
   class SwitchNode < Base
-    def initialize(base = nil, on_casters = [], else_caster = nil)
+    def initialize(base = nil, on_casters: [], else_caster: nil, pick_key: nil)
       @base = base
+      @pick_key = pick_key
 
       if Datacaster::Utils.pickable?(@base)
-        @base = Datacaster::Predefined.run { checked_key!(base) } &
-          Datacaster::Predefined.pick(base)
+        unless @pick_key.nil?
+          raise RuntimeError, "pick_key expected to be nil because #{@base.inspect} is pickable"
+        end
+        @pick_key = base
+        @base = Datacaster::Predefined.pick(base)
       end
 
       if !@base.nil? && !Datacaster.instance?(@base)
@@ -34,12 +38,12 @@ def on(caster_or_value, clause, strict: false)
 
       clause = DefinitionDSL.expand(clause)
 
-      self.class.new(@base, @ons + [[caster, clause]], @else)
+      self.class.new(@base, on_casters: @ons + [[caster, clause]], else_caster: @else, pick_key: @pick_key)
     end
 
     def else(else_caster)
       raise ArgumentError, "Datacaster: double else clause is not permitted", caller if @else
-      self.class.new(@base, @ons, else_caster)
+      self.class.new(@base, on_casters: @ons, else_caster: else_caster, pick_key: @pick_key)
     end
 
     def cast(object, runtime:)
@@ -59,7 +63,11 @@ def cast(object, runtime:)
         result = check.with_runtime(runtime).(switch_result)
         next unless result.valid?
 
-        return clause.with_runtime(runtime).(object)
+        if @pick_key.nil?
+          return clause.with_runtime(runtime).(object)
+        else
+          return runtime.checked_key!(@pick_key) { clause.with_runtime(runtime).(object) }
+        end
       end
 
       # all 'on'-s have failed
@@ -71,7 +79,7 @@ def cast(object, runtime:)
     end
 
     def inspect
-      "#<Datacaster::SwitchNode base: #{@base.inspect} on: #{@ons.inspect} else: #{@else.inspect}>"
+      "#<Datacaster::SwitchNode base: #{@base.inspect} on: #{@ons.inspect} else: #{@else.inspect} pick_key: #{@pick_key.inspect}>"
     end
   end
 end
diff --git a/lib/datacaster/version.rb b/lib/datacaster/version.rb
index 3072fb0..d65b611 100644
--- a/lib/datacaster/version.rb
+++ b/lib/datacaster/version.rb
@@ -1,3 +1,3 @@
 module Datacaster
-  VERSION = "3.2.2"
+  VERSION = "3.2.3"
 end
diff --git a/spec/switch_node_spec.rb b/spec/switch_node_spec.rb
index c114b01..060903a 100644
--- a/spec/switch_node_spec.rb
+++ b/spec/switch_node_spec.rb
@@ -136,5 +136,15 @@
 
       expect(caster.(kind: 1, name: '1').to_dry_result).to eq Success(kind: 1, name: '1')
     end
+
+    it 'marks matched-on value as checked with nested schemas' do
+      caster =
+        Datacaster.schema do
+          s = Datacaster.schema { string }
+          switch(:kind).on('person', hash_schema(kind: s))
+        end
+
+      expect(caster.(kind: 'person').to_dry_result).to eq Success(kind: 'person')
+    end
   end
 end