How to create custom-map whose get() call returns Optional wrapped value?

Question

I need a map that can hold null values and can return Optional wrapped value.

• I know HashMap can have null values, but it will not make it apparent to caller that value can be null.
• The easiest alternative is to use Map<K, Optional<V>>, but using Optional for data type is not ideal. (Relevant SO post about using Optional only for return: Use of Optional in a map)
• I couldn't extend HashMap class very well with its entrySet() method still facing the same null value issue, so I wrote my class that wraps HashMap and has get() and custom Entry.getValue() that returns Optional wrapped value. (still figuring out how to write a version of Collectors.toMap for this class)
• The problem now is the class is not easily "replaceable" with Map(not programming to interface) and a custom-not-much-extensible-map is circulating in my business logic that I am slightly uncomfortable with.

public class CustomMap<K, V> {
    private final Map<K, V> map;

    public CustomMap() {
        map = new HashMap<>();
    }

    public Optional<V> get(@NonNull final K k) {    // Lombok.NonNull
        return Optional.ofNullable(map.get(k));
    }

    public void put(@NonNull final K k, @Nullable final V v) {
        map.put(k, v);
    }

    public Set<Entry<K, V>> entrySet() {
        return map.entrySet().stream()
                .map(e -> new Entry(e.getKey(), e.getValue()))
                .collect(toImmutableSet());
    }

    public Set<K> keySet() {
        return map.keySet();
    }

    public static class Entry<K, V> {
        private final K k;
        private final V v;

        public Entry(K k, V v) {
            this.k = k;
            this.v = v;
        }

        public K getKey() {
            return k;
        }

        public Optional<V> getValue() {
            return Optional.ofNullable(v);
        }
    }
}

Expectation: customMap.get(K) should return Optional wrapped object and customMap.put(K, V) should not need to take Optional<V> as input. CustomMap should be able to act as Map.

There should be cleaner and extensible way to achieve this and I feel I am missing something obvious. Any suggestions?

Answer 1

If I understand you correctly, you want to make the use of Optionals apparent and you are reluctant to roll out some custom map implementation that deviates from the Map implementation; this conflicts with your statement that the Optional datatype is not ideal.

Assuming that you just do not want to store Optional instances in a map, a possible solution would be an Map implementation that wraps an (inner) Map<K,V> with null values as a Map<K, Optional<V>>, such as

import java.util.AbstractMap;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

public class OptionalMap<K, V> implements Map<K, Optional<V>> {
    private final Map<K, V> map;

    public OptionalMap(Map<K, V> map) {
        this.map = map;
    }

    @Override
    public void clear() {
        map.clear();
    }

    @Override
    public boolean containsKey(Object key) {
        return map.containsKey(key);
    }

    @Override
    public boolean containsValue(Object value) {
        if (!(value instanceof Optional)) {
            return false;
        }

        return map.containsValue(((Optional<?>)value).get());
    }

    @Override
    public Set<Entry<K, Optional<V>>> entrySet() {
        return
                map.entrySet().stream()
                .map(e -> new AbstractMap.SimpleEntry<K, Optional<V>>(e.getKey(), Optional.ofNullable(e.getValue())))
                .collect(Collectors.collectingAndThen(Collectors.toSet(), Collections::unmodifiableSet));
    }

    @Override
    public Optional<V> get(Object key) {
        return Optional.ofNullable(map.get(key));
    }

    @Override
    public boolean isEmpty() {
        return map.isEmpty();
    }

    @Override
    public Set<K> keySet() {
        return map.keySet();
    }

    @Override
    public Optional<V> put(K key, Optional<V> value) {
        final Optional<V> previous = Optional.ofNullable(map.get(key));
        map.put(key, value.orElse(null));
        return previous;
    }

    @Override
    public void putAll(Map<? extends K, ? extends Optional<V>> other) {
        Map<K, V> unwrappedMap =
            other.entrySet().stream()
            .<Entry<K, V>>map(entry -> new AbstractMap.SimpleEntry<K, V>(entry.getKey(), entry.getValue().orElse(null)))
            .collect(Collectors.toMap(Entry::getKey, Entry::getValue));

        map.putAll(unwrappedMap);
    }

    @Override
    public Optional<V> remove(Object key) {
        return Optional.ofNullable(map.remove(key));
    }

    @Override
    public int size() {
        return map.size();
    }

    @Override
    public Collection<Optional<V>> values() {
        return
                map.values().stream()
                .map(e -> Optional.ofNullable(e))
                .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
    }
}