WikidataFilter.java

package org.wikibrain.wikidata;

import org.wikibrain.core.lang.Language;

import java.util.*;

/**
 * @author Shilad Sen
 */
public class WikidataFilter {
    private Collection<Language> langs;

    // Entity types
    private Collection<WikidataEntity.Type> entityTypes;
    private Collection<Integer> entityIds;

    // For statements, but not entities
    private Collection<Integer> propertyIds;
    private Collection<WikidataStatement.Rank> ranks;
    private Collection<WikidataValue> values;

    public WikidataFilter() {
    }

    public Collection<Language> getLangs() {
        return langs;
    }

    public Collection<Short> getLangIds() {
        if (langs == null) {
            return null;
        }
        Set<Short> ids = new HashSet<Short>();
        for (Language lang : langs) {
            ids.add(lang.getId());
        }
        return ids;
    }

    public Collection<WikidataEntity.Type> getEntityTypes() {
        return entityTypes;
    }

    public Collection<String> getEntityTypeCodes() {
        if (entityTypes == null) {
            return null;
        }
        Set<String> typeCodes = new HashSet<String>();
        for (WikidataEntity.Type type : entityTypes) {
            typeCodes.add(Character.toString(type.code));
        }
        return typeCodes;
    }

    public Collection<WikidataValue> getValues() {
        return values;
    }

    public Collection<Integer> getEntityIds() {
        return entityIds;
    }

    public Collection<Integer> getPropertyIds() {
        return propertyIds;
    }

    public Collection<WikidataStatement.Rank> getRanks() {
        return ranks;
    }

    public Collection<Short> getRankOrdinals() {
        if (ranks == null) {
            return null;
        }
        Set<Short> ordinals = new HashSet<Short>();
        for (WikidataStatement.Rank rank : ranks) {
            ordinals.add((short)rank.ordinal());
        }
        return ordinals;
    }

    /**
     * Utility class to build a filter
     */
    public static class Builder {
        private WikidataFilter filter = new WikidataFilter();

        public Builder withLanguage(Language language) {
            filter.langs = Arrays.asList(language);
            return this;
        }

        public Builder withLanguages(Collection<Language> languages) {
            filter.langs = languages;
            return this;
        }

        public Builder withEntityType(WikidataEntity.Type type) {
            filter.entityTypes = Arrays.asList(type);
            return this;
        }

        public Builder withEntityType(Collection<WikidataEntity.Type> types) {
            filter.entityTypes = types;
            return this;
        }

        public Builder withEntityId(int id) {
            filter.entityIds = Arrays.asList(id);
            return this;
        }

        public Builder withEntityIds(Collection<Integer> ids) {
            filter.entityIds = ids;
            return this;
        }

        public Builder withPropertyId(int id) {
            filter.propertyIds = Arrays.asList(id);
            return this;
        }

        public Builder withPropertyIds(Collection<Integer> ids) {
            filter.propertyIds = ids;
            return this;
        }

        public Builder withRank(WikidataStatement.Rank rank) {
            filter.ranks = Arrays.asList(rank);
            return this;
        }

        public Builder withRanks(Collection<WikidataStatement.Rank> ranks) {
            filter.ranks = ranks;
            return this;
        }

        public Builder withValue(WikidataValue value) {
            filter.values = Arrays.asList(value);
            return this;
        }

        public Builder withValues(Collection<WikidataValue> values) {
            filter.values = values;
            return this;
        }

        public WikidataFilter build() {
            return filter;
        }
    }
}