Thread#fetch

* thread.c (rb_thread_fetch): add new method Thread#fetch.
  [Feature #13009]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@57683 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

NEWS

@@ -24,6 +24,10 @@ with all sufficient information, see the ChangeLog file or Redmine
   * Update Onigmo 6.1.1.
     * Support absent operator https://github.com/k-takata/Onigmo/issues/82
 
+* Thread
+
+  * Thread#fetch  [Feature #13009]
+
 === Stdlib updates (outstanding ones only)
 
 === Compatibility issues (excluding feature bug fixes)

test/ruby/test_thread.rb

@@ -482,6 +482,36 @@ class TestThread < Test::Unit::TestCase
     t.kill if t
   end
 
+  def test_thread_local_fetch
+    t = Thread.new { sleep }
+
+    assert_equal(false, t.key?(:foo))
+
+    t["foo"] = "foo"
+    t["bar"] = "bar"
+    t["baz"] = "baz"
+
+    x = nil
+    assert_equal("foo", t.fetch(:foo, 0))
+    assert_equal("foo", t.fetch(:foo) {x = true})
+    assert_nil(x)
+    assert_equal("foo", t.fetch("foo", 0))
+    assert_equal("foo", t.fetch("foo") {x = true})
+    assert_nil(x)
+
+    x = nil
+    assert_equal(0, t.fetch(:qux, 0))
+    assert_equal(1, t.fetch(:qux) {x = 1})
+    assert_equal(1, x)
+    assert_equal(2, t.fetch("qux", 2))
+    assert_equal(3, t.fetch("qux") {x = 3})
+    assert_equal(3, x)
+
+    assert_raise(KeyError) {t.fetch(:qux)}
+  ensure
+    t.kill if t
+  end
+
   def test_thread_local_security
     assert_raise(RuntimeError) do
       Thread.new do

thread.c

@@ -3090,6 +3090,39 @@ rb_thread_aref(VALUE thread, VALUE key)
     return rb_thread_local_aref(thread, id);
 }
 
+static VALUE
+rb_thread_fetch(int argc, VALUE *argv, VALUE self)
+{
+    VALUE key, val;
+    ID id;
+    rb_thread_t *th;
+    int block_given;
+
+    rb_check_arity(argc, 1, 2);
+    key = argv[0];
+
+    block_given = rb_block_given_p();
+    if (block_given && argc == 2) {
+	rb_warn("block supersedes default value argument");
+    }
+
+    id = rb_check_id(&key);
+    GetThreadPtr(self, th);
+
+    if (id == recursive_key) {
+	return th->local_storage_recursive_hash;
+    }
+    if (id && th->local_storage && st_lookup(th->local_storage, id, &val)) {
+	return val;
+    }
+    if (block_given)
+	return rb_yield(key);
+    else if (argc == 1)
+	rb_raise(rb_eKeyError, "key not found: %"PRIsVALUE, key);
+    else
+	return argv[1];
+}
+
 static VALUE
 threadptr_local_aset(rb_thread_t *th, ID id, VALUE val)
 {
@@ -4796,6 +4829,7 @@ Init_Thread(void)
     rb_define_method(rb_cThread, "wakeup", rb_thread_wakeup, 0);
     rb_define_method(rb_cThread, "[]", rb_thread_aref, 1);
     rb_define_method(rb_cThread, "[]=", rb_thread_aset, 2);
+    rb_define_method(rb_cThread, "fetch", rb_thread_fetch, -1);
     rb_define_method(rb_cThread, "key?", rb_thread_key_p, 1);
     rb_define_method(rb_cThread, "keys", rb_thread_keys, 0);
     rb_define_method(rb_cThread, "priority", rb_thread_priority, 0);