diff --git a/src/dicts.jl b/src/dicts.jl index 336ae4a..eb2b5ae 100644 --- a/src/dicts.jl +++ b/src/dicts.jl @@ -117,6 +117,58 @@ function Base.get(d::LMDBDict{K,V}, key, default) where {K,V} end end +function Base.get!(d::LMDBDict{K,V}, key, default) where {K,V} + txn_dbi_do(d, readonly = true) do txn, dbi + mdb_key_ref = Ref(MDBValue(toref(convert(K,key)))) + mdb_val_ref = Ref(MDBValue()) + # Get value + ret = _mdb_get(txn.handle, dbi.handle, mdb_key_ref, mdb_val_ref) + if ret == MDB_NOTFOUND + d[key] = default + return default + elseif ret == Cint(0) + return mbd_unpack(V, mdb_val_ref) + else + throw(LMDB.LMDBError(ret)) + end + end +end + +function Base.get(f::F, d::LMDBDict{K,V}, key) where {K,V,F<:Union{Function, Type}} + txn_dbi_do(d, readonly = true) do txn, dbi + mdb_key_ref = Ref(MDBValue(toref(convert(K,key)))) + mdb_val_ref = Ref(MDBValue()) + # Get value + ret = _mdb_get(txn.handle, dbi.handle, mdb_key_ref, mdb_val_ref) + if ret == MDB_NOTFOUND + default = f() + return default + elseif ret == Cint(0) + return mbd_unpack(V, mdb_val_ref) + else + throw(LMDB.LMDBError(ret)) + end + end +end + +function Base.get!(f::F, d::LMDBDict{K,V}, key) where {K,V,F<:Union{Function, Type}} + txn_dbi_do(d, readonly = true) do txn, dbi + mdb_key_ref = Ref(MDBValue(toref(convert(K,key)))) + mdb_val_ref = Ref(MDBValue()) + # Get value + ret = _mdb_get(txn.handle, dbi.handle, mdb_key_ref, mdb_val_ref) + if ret == MDB_NOTFOUND + default = f() + d[key] = default + return default + elseif ret == Cint(0) + return mbd_unpack(V, mdb_val_ref) + else + throw(LMDB.LMDBError(ret)) + end + end +end + function Base.setindex!(d::LMDBDict{K,V},v,k) where {K,V} txn_dbi_do(d) do txn, dbi LMDB.put!(txn,dbi,convert(K,k),convert(V,v)) diff --git a/test/dict.jl b/test/dict.jl index e6c36aa..dc6bd07 100644 --- a/test/dict.jl +++ b/test/dict.jl @@ -49,4 +49,34 @@ d["b"] = [0,0,0] @test LMDB.list_dirs(d) == ["aa/", "b"] @test LMDB.list_dirs(d,prefix="aa/") == ["aa/a", "aa/b", "aa/c"] @test LMDB.valuesize(d,prefix="aa/") == sizeof(Float32)*18 -end \ No newline at end of file + +@testset "Tests for get and get!" begin + mktempdir() do dir + d = LMDBDict{String, String}(dir) + @test !haskey(d, "foo") + @test get(d, "foo", "bar") == "bar" + @test !haskey(d, "foo") + @test get!(d, "foo", "bar") == "bar" + @test haskey(d, "foo") + @test d["foo"] == "bar" + @test get(d, "foo", "hello") == "bar" + @test d["foo"] == "bar" + @test get!(d, "foo", "hello") == "bar" + @test d["foo"] == "bar" + end + mktempdir() do dir + d = LMDBDict{String, String}(dir) + @test !haskey(d, "foo") + @test get(() -> "bar", d, "foo") == "bar" + @test !haskey(d, "foo") + @test get!(() -> "bar", d, "foo") == "bar" + @test haskey(d, "foo") + @test d["foo"] == "bar" + @test get(() -> "hello", d, "foo") == "bar" + @test d["foo"] == "bar" + @test get!(() -> "hello", d, "foo") == "bar" + @test d["foo"] == "bar" + end +end + +end