add numbertype=Float64 option (#240)

quinnj · Nov 3, 2022 · 765c3bf · 765c3bf
1 parent df5049d
commit 765c3bf
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 11 deletions.
diff --git a/src/read.jl b/src/read.jl
@@ -25,8 +25,11 @@ Read JSON.
 * `allow_inf`: Allow reading of `Inf` and `NaN` values (not part of the JSON standard). [default `false`]
 * `dateformat`: A [`DateFormat`](https://docs.julialang.org/en/v1/stdlib/Dates/#Dates.DateFormat) describing the format of dates in the JSON so that they can be read into `Date`s, `Time`s, or `DateTime`s when reading into a type. [default `Dates.default_format(T)`]
 * `parsequoted`: Accept quoted values when reading into a NumberType. [default `false`]
+* `numbertype`: Type to parse numbers as. [default `nothing`, which parses numbers as Int if possible, Float64 otherwise]
 """
-function read(str::AbstractString; jsonlines::Bool=false, kw...)
+function read(str::AbstractString; jsonlines::Bool=false,
+              numbertype::Union{DataType, Nothing}=nothing, kw...)
+
     buf = codeunits(str)
     len = length(buf)
     if len == 0
@@ -39,10 +42,18 @@ function read(str::AbstractString; jsonlines::Bool=false, kw...)
     @wh
     tape = len < 1000 ? Vector{UInt64}(undef, len + 4) :
         Vector{UInt64}(undef, div(len, 10))
+    if numbertype === nothing
+        checkint = true
+    elseif numbertype == Float64
+        checkint = false
+    else
+        throw(ArgumentError("numbertype $numbertype is not supported. " *
+                            "Only `nothing` (default) and `Float64` are supported so far."))
+    end
     if jsonlines
-        pos, tapeidx = jsonlines!(buf, pos, len, b, tape, Int64(1); kw...)
+        pos, tapeidx = jsonlines!(buf, pos, len, b, tape, Int64(1), checkint; kw...)
     else
-        pos, tapeidx = read!(buf, pos, len, b, tape, Int64(1), Any; kw...)
+        pos, tapeidx = read!(buf, pos, len, b, tape, Int64(1), Any, checkint; kw...)
     end
     @inbounds t = tape[1]
     if isobject(t)
@@ -74,9 +85,9 @@ const FLOAT_INT_BOUND = 2.0^53
 
 function read!(buf, pos, len, b, tape, tapeidx, ::Type{Any}, checkint=true; allow_inf::Bool=false)
     if b == UInt8('{')
-        return read!(buf, pos, len, b, tape, tapeidx, Object; allow_inf=allow_inf)
+        return read!(buf, pos, len, b, tape, tapeidx, Object, checkint; allow_inf=allow_inf)
     elseif b == UInt8('[')
-        return read!(buf, pos, len, b, tape, tapeidx, Array; allow_inf=allow_inf)
+        return read!(buf, pos, len, b, tape, tapeidx, Array, checkint; allow_inf=allow_inf)
     elseif b == UInt8('"')
         return read!(buf, pos, len, b, tape, tapeidx, String)
     elseif b == UInt8('n')
@@ -203,7 +214,7 @@ function read!(buf, pos, len, b, tape, tapeidx, ::Type{Nothing})
     end
 end
 
-function read!(buf, pos, len, b, tape, tapeidx, ::Type{Object}; kw...)
+function read!(buf, pos, len, b, tape, tapeidx, ::Type{Object}, checkint=true; kw...)
     objidx = tapeidx
     eT = EMPTY
     pos += 1
@@ -263,7 +274,7 @@ function read!(buf, pos, len, b, tape, tapeidx, ::Type{Object}; kw...)
         @wh
         # now positioned at start of value
         prevtapeidx = tapeidx
-        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any; kw...)
+        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any, checkint; kw...)
         @eof
         b = getbyte(buf, pos)
         @wh
@@ -297,7 +308,7 @@ function read!(buf, pos, len, b, tape, tapeidx, ::Type{Object}; kw...)
     invalid(error, buf, pos, Object)
 end
 
-function read!(buf, pos, len, b, tape, tapeidx, ::Type{Array}; kw...)
+function read!(buf, pos, len, b, tape, tapeidx, ::Type{Array}, checkint=true; kw...)
     arridx = tapeidx
     eT = EMPTY
     pos += 1
@@ -317,7 +328,8 @@ function read!(buf, pos, len, b, tape, tapeidx, ::Type{Array}; kw...)
     while true
         # positioned at start of value
         prevtapeidx = tapeidx
-        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any, eT != FLOAT && eT != (FLOAT | NULL); kw...)
+        check_int = checkint ? eT != FLOAT && eT != (FLOAT | NULL) : false
+        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any, check_int; kw...)
         @eof
         b = getbyte(buf, pos)
         @wh
@@ -345,7 +357,7 @@ function read!(buf, pos, len, b, tape, tapeidx, ::Type{Array}; kw...)
     invalid(error, buf, pos, Array)
 end
 
-function jsonlines!(buf, pos, len, b, tape, tapeidx; kw...)
+function jsonlines!(buf, pos, len, b, tape, tapeidx, checkint=true; kw...)
     arridx = tapeidx
     eT = EMPTY
     if pos > len
@@ -361,7 +373,8 @@ function jsonlines!(buf, pos, len, b, tape, tapeidx; kw...)
         @wh
         # positioned at start of value
         prevtapeidx = tapeidx
-        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any, eT != FLOAT && eT != (FLOAT | NULL); kw...)
+        check_int = checkint ? eT != FLOAT && eT != (FLOAT | NULL) : false
+        pos, tapeidx = read!(buf, pos, len, b, tape, tapeidx, Any, check_int; kw...)
         @inbounds eT = promoteeltype(eT, gettypemask(tape[prevtapeidx]))
         nelem += 1
         if pos > len

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -172,6 +172,15 @@ roundtrip(x::T) where T = JSON3.read(JSON3.write(x), T)
 @test JSON3.read("true") === true
 @test JSON3.read("false") === false
 
+# numbertype
+@test JSON3.read("1.0", numbertype=Float64) === 1.0
+@test JSON3.read("1", numbertype=Float64) === 1.0
+@test_throws ArgumentError JSON3.read("1", numbertype=Float32)
+@test JSON3.read("{\"hey\":1}", numbertype=Float64).hey === 1.0
+@test JSON3.read("[\"hey\",1]", numbertype=Float64)[2] === 1.0
+@test JSON3.read("[[1]]", numbertype=Float64)[1][1] === 1.0
+@test JSON3.read("1.0"; jsonlines=true, numbertype=Float64)[1] === 1.0
+
 # more robust Int vs. Float detection: https://github.com/quinnj/JSON3.jl/issues/124
 @test JSON3.read("9223372036854775807") == 9223372036854775807
 @test JSON3.read("9223372036854775808") == 9.223372036854776e18 # overflowed Int64 promotes to Float64 (though slightly lossy)