Merge 1c6b039 into 94f4542

wildart · Sep 1, 2022 · 5827887 · 5827887
2 parents 94f4542 + 1c6b039
commit 5827887
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 1 deletion.
diff --git 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
@@ -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
+
+@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