JSON::Coder callback now recieve a second argument to mark object keys

byroot · byroot · commit 7982b6148cba · 2025-09-18T21:31:36.000+02:00
e.g.
```ruby
{ 1 =&gt; 2 }
```

The callback will be invoked for `1` as while it has a native JSON
equivalent, it's not legal as an object name.
diff --git a/CHANGES.md b/CHANGES.md
@@ -2,6 +2,7 @@
 
 ### Unreleased
 
+* `JSON::Coder` callback now receive a second argument to convey whether the object is a hash key.
 * Tuned the floating point number generator to not use scientific notation as agressively.
 
 ### 2025-09-18 (2.14.1)
diff --git a/README.md b/README.md
@@ -97,7 +97,7 @@ Instead it is recommended to use the newer `JSON::Coder` API:
 
 ```ruby
 module MyApp
-  API_JSON_CODER = JSON::Coder.new do |object|
+  API_JSON_CODER = JSON::Coder.new do |object, is_object_key|
     case object
     when Time
       object.iso8601(3)
@@ -113,6 +113,8 @@ puts MyApp::API_JSON_CODER.dump(Time.now.utc) # => "2025-01-21T08:41:44.286Z"
 The provided block is called for all objects that don't have a native JSON equivalent, and
 must return a Ruby object that has a native JSON equivalent.
 
+It is also called for objects that do have a JSON equivalent, but are used as Hash keys, for instance `{ 1 => 2}`.
+
 ## Combining JSON fragments
 
 To combine JSON fragments into a bigger JSON document, you can use `JSON::Fragment`:
diff --git a/ext/json/ext/generator/generator.c b/ext/json/ext/generator/generator.c
@@ -29,6 +29,7 @@ typedef struct JSON_Generator_StateStruct {
 
     enum duplicate_key_action on_duplicate_key;
 
+    bool as_json_single_arg;
     bool allow_nan;
     bool ascii_only;
     bool script_safe;
@@ -1033,6 +1034,13 @@ json_inspect_hash_with_mixed_keys(struct hash_foreach_arg *arg)
     }
 }
 
+static VALUE
+json_call_as_json(JSON_Generator_State *state, VALUE object, VALUE is_key)
+{
+    VALUE proc_args[2] = {object, is_key};
+    return rb_proc_call_with_block(state->as_json, 2, proc_args, Qnil);
+}
+
 static int
 json_object_i(VALUE key, VALUE val, VALUE _arg)
 {
@@ -1086,7 +1094,7 @@ json_object_i(VALUE key, VALUE val, VALUE _arg)
         default:
             if (data->state->strict) {
                 if (RTEST(data->state->as_json) && !as_json_called) {
-                    key = rb_proc_call_with_block(data->state->as_json, 1, &key, Qnil);
+                    key = json_call_as_json(data->state, key, Qtrue);
                     key_type = rb_type(key);
                     as_json_called = true;
                     goto start;
@@ -1328,7 +1336,7 @@ static void generate_json_float(FBuffer *buffer, struct generate_json_data *data
         /* for NaN and Infinity values we either raise an error or rely on Float#to_s. */
         if (!allow_nan) {
             if (data->state->strict && data->state->as_json) {
-                VALUE casted_obj = rb_proc_call_with_block(data->state->as_json, 1, &obj, Qnil);
+                VALUE casted_obj = json_call_as_json(data->state, obj, Qfalse);
                 if (casted_obj != obj) {
                     increase_depth(data);
                     generate_json(buffer, data, casted_obj);
@@ -1416,7 +1424,7 @@ static void generate_json(FBuffer *buffer, struct generate_json_data *data, VALU
             general:
                 if (data->state->strict) {
                     if (RTEST(data->state->as_json) && !as_json_called) {
-                        obj = rb_proc_call_with_block(data->state->as_json, 1, &obj, Qnil);
+                        obj = json_call_as_json(data->state, obj, Qfalse);
                         as_json_called = true;
                         goto start;
                     } else {
@@ -1942,6 +1950,7 @@ static int configure_state_i(VALUE key, VALUE val, VALUE _arg)
     else if (key == sym_allow_duplicate_key)   { state->on_duplicate_key = RTEST(val) ? JSON_IGNORE : JSON_RAISE; }
     else if (key == sym_as_json)               {
         VALUE proc = RTEST(val) ? rb_convert_type(val, T_DATA, "Proc", "to_proc") : Qfalse;
+        state->as_json_single_arg = proc && rb_proc_arity(proc) == 1;
         state_write_value(data, &state->as_json, proc);
     }
     return ST_CONTINUE;
diff --git a/java/src/json/ext/Generator.java b/java/src/json/ext/Generator.java
@@ -396,7 +396,7 @@ static void generateFloat(ThreadContext context, Session session, RubyFloat obje
 
             if (!state.allowNaN()) {
                 if (state.strict() && state.getAsJSON() != null) {
-                    IRubyObject castedValue = state.getAsJSON().call(context, object);
+                    IRubyObject castedValue = state.getAsJSON().call(context, object, context.getRuntime().getFalse());
                     if (castedValue != object) {
                         getHandlerFor(context.runtime, castedValue).generate(context, session, castedValue, buffer);
                         return;
@@ -623,7 +623,7 @@ private static void processEntry(ThreadContext context, Session session, OutputS
                 GeneratorState state = session.getState(context);
                 if (state.strict()) {
                     if (state.getAsJSON() != null) {
-                        key = state.getAsJSON().call(context, key);
+                        key = state.getAsJSON().call(context, key, context.getRuntime().getTrue());
                         keyStr = castKey(context, key);
                     }
 
@@ -760,7 +760,7 @@ static RubyString generateGenericNew(ThreadContext context, Session session, IRu
         GeneratorState state = session.getState(context);
         if (state.strict()) {
             if (state.getAsJSON() != null) {
-                IRubyObject value = state.getAsJSON().call(context, object);
+                IRubyObject value = state.getAsJSON().call(context, object, context.getRuntime().getFalse());
                 Handler handler = getHandlerFor(context.runtime, value);
                 if (handler == GENERIC_HANDLER) {
                     throw Utils.buildGeneratorError(context, object, value + " returned by as_json not allowed in JSON").toThrowable();
diff --git a/lib/json/truffle_ruby/generator.rb b/lib/json/truffle_ruby/generator.rb
@@ -450,7 +450,7 @@ def to_json(state = nil, *)
               value = self
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value)
                 if state.as_json
-                  value = state.as_json.call(value)
+                  value = state.as_json.call(value, false)
                   unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value
                     raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
                   end
@@ -511,7 +511,7 @@ def json_transform(state)
 
               if state.strict? && !(Symbol === key || String === key)
                 if state.as_json
-                  key = state.as_json.call(key)
+                  key = state.as_json.call(key, true)
                 end
 
                 unless Symbol === key || String === key
@@ -529,7 +529,7 @@ def json_transform(state)
               result = +"#{result}#{key_json}#{state.space_before}:#{state.space}"
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value)
                 if state.as_json
-                  value = state.as_json.call(value)
+                  value = state.as_json.call(value, false)
                   unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value
                     raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
                   end
@@ -590,7 +590,7 @@ def json_transform(state)
               result << state.indent * depth if indent
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value || Symbol == value)
                 if state.as_json
-                  value = state.as_json.call(value)
+                  value = state.as_json.call(value, false)
                   unless false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value || Fragment === value || Symbol === value
                     raise GeneratorError.new("#{value.class} returned by #{state.as_json} not allowed in JSON", value)
                   end
@@ -625,7 +625,7 @@ def to_json(state = nil, *args)
               if state.allow_nan?
                 to_s
               elsif state.strict? && state.as_json
-                casted_value = state.as_json.call(self)
+                casted_value = state.as_json.call(self, false)
 
                 if casted_value.equal?(self)
                   raise GeneratorError.new("#{self} not allowed in JSON", self)
diff --git a/test/json/json_coder_test.rb b/test/json/json_coder_test.rb
@@ -12,15 +12,19 @@ def test_json_coder_with_proc
   end
 
   def test_json_coder_with_proc_with_unsupported_value
-    coder = JSON::Coder.new do |object|
+    coder = JSON::Coder.new do |object, is_key|
+      assert_equal false, is_key
       Object.new
     end
     assert_raise(JSON::GeneratorError) { coder.dump([Object.new]) }
   end
 
   def test_json_coder_hash_key
     obj = Object.new
-    coder = JSON::Coder.new(&:to_s)
+    coder = JSON::Coder.new do |obj, is_key|
+      assert_equal true, is_key
+      obj.to_s
+    end
     assert_equal %({#{obj.to_s.inspect}:1}), coder.dump({ obj => 1 })
 
     coder = JSON::Coder.new { 42 }
@@ -49,14 +53,14 @@ def test_json_coder_load_options
   end
 
   def test_json_coder_dump_NaN_or_Infinity
-    coder = JSON::Coder.new(&:inspect)
+    coder = JSON::Coder.new { |o| o.inspect }
     assert_equal "NaN", coder.load(coder.dump(Float::NAN))
     assert_equal "Infinity", coder.load(coder.dump(Float::INFINITY))
     assert_equal "-Infinity", coder.load(coder.dump(-Float::INFINITY))
   end
 
   def test_json_coder_dump_NaN_or_Infinity_loop
-    coder = JSON::Coder.new(&:itself)
+    coder = JSON::Coder.new { |o| o.itself }
     error = assert_raise JSON::GeneratorError do
       coder.dump(Float::NAN)
     end
diff --git a/test/json/json_generator_test.rb b/test/json/json_generator_test.rb
@@ -822,7 +822,7 @@ def test_fragment
 
   def test_json_generate_as_json_convert_to_proc
     object = Object.new
-    assert_equal object.object_id.to_json, JSON.generate(object, strict: true, as_json: :object_id)
+    assert_equal object.object_id.to_json, JSON.generate(object, strict: true, as_json: -> (o, is_key) { o.object_id })
   end
 
   def assert_float_roundtrip(expected, actual)
