Storing multiple object types in Java maps with type safety

Java’s generics have undoubtedly made life easier in some cases by enforcing type safety on collections. One collection that frequently does not benefit from them is maps used in a heterogeneous way – that is, a java.util.Map (or equivalent) that contains several non-compatible values. Assuming the keys are of the same type, it’s still not possible to type the value of the map as anything other than a java.lang.Object – exactly what it would have been prior to generics.

Looking through the Jetbrains documentation on its API for IntelliJ plugin development, it’s interesting to note their workaround for enforcing type safety in a map containing heterogeneous values:

• DataKey<T> – the access point to the map
• DataContext – an interface that can be used to wrap a map

This approach externalises the type-safety of the map’s content by shifting the casts to DataKey.

The interesting part of DataKey, from this viewpoint, is

@Nullable
public T getData(@NotNull DataContext dataContext)
{
    return (T) dataContext.getData(myName);
}

Read access is via the key, and not directly on the map itself, e.g. for a key

DataKey<Foo> FOO = DataKey.create("foo")

a type-safe call without a cast can be made on the wrapped via:

Foo foo = MyKeys.FOO.getData(dataContext);

The map-based implementations of DataContext just does a regular lookup with the key provided by DataKey:

public class MyDataContext implements DataContext
{
    private final Map<String, ?> data = new HashMap<String, ?>();
    @Nullable Object getData(@NonNls String dataId)
    {
        return data.get(dataId);
    }
}

and the returned object is cast to the correct type on the way out of DataKey. Simple and elegant.

It’s worth noting that both DataKey and DataContext are read-only in function – there’s a getData(), but no setData().   It’s easy enough to extend the concept to make sure that what goes into the map is of the same type as what comes out by adding a couple of extra methods, one on DataKey:

    public void setData(DataContext dataContext, T value)
    {
        dataContext.setData(getName(), value);
    }

and another on DataContext:

    public void setData(String key, Object value)
    {
        data.put(key, value);
    }

I’ve seen a lot of code where keys to access values from maps have been declared as static final Strings. Using DataKey, you could replace these declarations with DataKeys that specify the key to access the value and the type of the value itself, so

public static final String USER = "user";
...
User user = (User)map.get(USER);

would become

public static final DataKey<User> USER = DataKey.create("user");
...
User user = USER.getData(dataContext);