Add AbstractDict interface

docs/src/index.md

@@ -39,9 +39,18 @@ For multiple elements you can similarly use
 pseudofile.(family, [:C, :Si])
 ```
 
-Some metadata information is stored in the `family` object:
+A `PseudoFamily` is furthermore an `AbstractDict{Symbol,String}`
+for the mapping of element symbol to file path, e.g. one can perform
+index lookup
 ```@example index-example
-family
+family[:Si]
+```
+or iterate over pairs
+```@example index-example
+for (k, v) in family
+   println(k, " => ", v)
+   break
+end
 ```
 
 For a list of available identifiers see

src/PseudoPotentialData.jl

@@ -8,6 +8,18 @@ export PseudoFamily, pseudofile
 @compat public families
 @compat public family_identifiers
 
+export available_elements, has_element
+
 include("pseudofamily.jl")
 
+"""Get the list of available pseudopotential family identifiers."""
+function family_identifiers()
+    artifact_file = find_artifacts_toml(@__FILE__)
+    @assert !isnothing(artifact_file)
+    collect(keys(TOML.parsefile(artifact_file)))
+end
+
+"""The list of all known pseudopotential families."""
+const families = map(PseudoFamily, family_identifiers())
+
 end

src/pseudofamily.jl

@@ -1,5 +1,6 @@
-struct PseudoFamily
+struct PseudoFamily <: AbstractDict{Symbol,String}
     identifier::String
+    #
     # metadata
     extension::String   # Filename expected as $(symbol).$(extension)
     functional::String  # DFT functional keyword (or "" if unspecified)
@@ -24,38 +25,61 @@ function PseudoFamily(identifier::AbstractString)
                  meta["functional"],
                  VersionNumber(meta["version"]))
 end
-
 Base.Broadcast.broadcastable(l::PseudoFamily) = Ref(l)
-Base.getindex(family::PseudoFamily, element::Symbol) = pseudofile(family, element)
-
 
-"""Get the list of available pseudopotential family identifiers."""
-function family_identifiers()
-    artifact_file = find_artifacts_toml(@__FILE__)
-    @assert !isnothing(artifact_file)
-    collect(keys(TOML.parsefile(artifact_file)))
+function Base.show(io::IO, family::PseudoFamily)
+    print(io, "PseudoFamily(\"$(family.identifier)\")")
+end
+function Base.show(io::IO, ::MIME"text/plain", family::PseudoFamily)
+    show(io, family)
 end
 
-"""The list of all known pseudopotential families."""
-const families = map(PseudoFamily, family_identifiers())
-
+#
+# Helper functions (not exported)
+#
 
 """
 Return the directory containing the pseudo files.
 This downloads the artifact if necessary.
 """
 artifact_directory(family::PseudoFamily) = (@artifact_str "$(family.identifier)")
 
+"""Return the list of all pseudopotential files in the artifact"""
+function available_elements(family::PseudoFamily)
+    # TODO Once this is part of the metadata, do this without downloading
+    files = filter!(endswith(family.extension),
+                    readdir(artifact_directory(family)))
+    map(files) do file
+        base, _ = splitext(file)
+        Symbol(base)
+    end
+end
+
 """
 Get the full path to the file containing the pseudopotential information
-for a particular element and a particular pseudopotential `family`.
-The family can be specified as an identifier or an object.
+for a particular `element` (identified by an atomic symbol) and a particular
+pseudopotential `family`.
 """
-function pseudofile(family::PseudoFamily, element::Symbol)
+pseudofile(family::PseudoFamily, element::Symbol) = family[element]
+
+#
+# AbstractDict interface
+#
+
+Base.keys(family::PseudoFamily) = available_elements(family)
+Base.length(family::PseudoFamily) = length(available_elements(family))
+
+function Base.getindex(family::PseudoFamily, element::Symbol)
     file = joinpath(artifact_directory(family), "$(element)." * family.extension)
     isfile(file) || throw(KeyError(element))
     file
 end
-function pseudofile(family::AbstractString, element::Symbol)
-    pseudofile(PseudoFamily(family), element)
+
+function Base.iterate(family::PseudoFamily, state::Int=0)
+    if state == length(family)
+        return nothing
+    else
+        element = available_elements(family)[state+1]
+        return (element => family[element], state+1)
+    end
 end

test/runtests.jl

@@ -10,17 +10,15 @@ using Test
         @test family.functional == "lda"
         @test family.version == v"0.4.1"
 
-        file = pseudofile("pd_nc_sr_lda_stringent_0.4.1_upf", :Si)
+        file = pseudofile(family, :Si)
         @test basename(file) == "Si.upf"
-        @test file == pseudofile(family, :Si)
         @test file == family[:Si]
 
         files = pseudofile.(family, [:Si, :Si])
         @test all(isequal(file), files)
 
         @test_throws KeyError pseudofile(family, :Uun)
         @test_throws KeyError family[:Uun]
-        @test_throws KeyError pseudofile(identifier, :Uub)
     end
 
     @testset "Test all libraries can be loaded" begin
@@ -29,4 +27,18 @@ using Test
             @test family.identifier == identifier
         end
     end
+
+    @testset "Dict interface of PseudoFamily" begin
+        identifier = "pd_nc_sr_pbe_stringent_0.4.1_upf"
+        family = PseudoFamily(identifier)
+        @test length(family) == 72
+        @test   :Si  in keys(family)
+        @test !(:Uub in keys(family))
+        @test  haskey(family, :Si)
+        @test !haskey(family, :Uub)
+
+        basedir = PseudoPotentialData.artifact_directory(family)
+        @test joinpath(basedir, "Si.upf") in values(family)
+        @test family[:Si] == joinpath(basedir, "Si.upf")
+    end
 end