Add entity types to Options.

Test: bit FrameworksCoreTests:android.view.textclassifier.TextClassificationManagerTest
Bug: b/67629726
Change-Id: I9cad0159ab539a71d9f504019ebe91fe18206d60

3 files changed, 150 insertions, 6 deletions

diff --git a/core/java/android/view/textclassifier/TextClassifier.java b/core/java/android/view/textclassifier/TextClassifier.java
index fdc9f92347db..ed6043038600 100644
--- a/core/java/android/view/textclassifier/TextClassifier.java
+++ b/core/java/android/view/textclassifier/TextClassifier.java
@@ -16,17 +16,23 @@
 
 package android.view.textclassifier;
 
+import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.StringDef;
 import android.annotation.WorkerThread;
 import android.os.LocaleList;
+import android.util.ArraySet;
 
 import com.android.internal.util.Preconditions;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 
 /**
  * Interface for providing text classification related features.
@@ -58,6 +64,20 @@ public interface TextClassifier {
     })
     @interface EntityType {}
 
+    /** Designates that the TextClassifier should identify all entity types it can. **/
+    int ENTITY_PRESET_ALL = 0;
+    /** Designates that the TextClassifier should identify no entities. **/
+    int ENTITY_PRESET_NONE = 1;
+    /** Designates that the TextClassifier should identify a base set of entities determined by the
+     * TextClassifier. **/
+    int ENTITY_PRESET_BASE = 2;
+
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(prefix = { "ENTITY_CONFIG_" },
+            value = {ENTITY_PRESET_ALL, ENTITY_PRESET_NONE, ENTITY_PRESET_BASE})
+    @interface EntityPreset {}
+
     /**
      * No-op TextClassifier.
      * This may be used to turn off TextClassifier features.
@@ -217,6 +237,8 @@ public interface TextClassifier {
      * Returns a {@link TextLinks} that may be applied to the text to annotate it with links
      * information.
      *
+     * If no options are supplied, default values will be used, determined by the TextClassifier.
+     *
      * @param text the text to generate annotations for
      * @param options configuration for link generation
      *
@@ -251,6 +273,16 @@ public interface TextClassifier {
     }
 
     /**
+     * Returns a {@link Collection} of the entity types in the specified preset.
+     *
+     * @see #ENTITIES_ALL
+     * @see #ENTITIES_NONE
+     */
+    default Collection<String> getEntitiesForPreset(@EntityPreset int entityPreset) {
+        return Collections.EMPTY_LIST;
+    }
+
+    /**
      * Logs a TextClassifier event.
      *
      * @param source the text classifier used to generate this event
@@ -268,6 +300,62 @@ public interface TextClassifier {
         return TextClassifierConstants.DEFAULT;
     }
 
+    /**
+     * Configuration object for specifying what entities to identify.
+     *
+     * Configs are initially based on a predefined preset, and can be modified from there.
+     */
+    final class EntityConfig {
+        private final @TextClassifier.EntityPreset int mEntityPreset;
+        private final Collection<String> mExcludedEntityTypes;
+        private final Collection<String> mIncludedEntityTypes;
+
+        public EntityConfig(@TextClassifier.EntityPreset int mEntityPreset) {
+            this.mEntityPreset = mEntityPreset;
+            mExcludedEntityTypes = new ArraySet<>();
+            mIncludedEntityTypes = new ArraySet<>();
+        }
+
+        /**
+         * Specifies an entity to include in addition to any specified by the enity preset.
+         *
+         * Note that if an entity has been excluded, the exclusion will take precedence.
+         */
+        public EntityConfig includeEntities(String... entities) {
+            for (String entity : entities) {
+                mIncludedEntityTypes.add(entity);
+            }
+            return this;
+        }
+
+        /**
+         * Specifies an entity to be excluded.
+         */
+        public EntityConfig excludeEntities(String... entities) {
+            for (String entity : entities) {
+                mExcludedEntityTypes.add(entity);
+            }
+            return this;
+        }
+
+        /**
+         * Returns an unmodifiable list of the final set of entities to find.
+         */
+        public List<String> getEntities(TextClassifier textClassifier) {
+            ArrayList<String> entities = new ArrayList<>();
+            for (String entity : textClassifier.getEntitiesForPreset(mEntityPreset)) {
+                if (!mExcludedEntityTypes.contains(entity)) {
+                    entities.add(entity);
+                }
+            }
+            for (String entity : mIncludedEntityTypes) {
+                if (!mExcludedEntityTypes.contains(entity) && !entities.contains(entity)) {
+                    entities.add(entity);
+                }
+            }
+            return Collections.unmodifiableList(entities);
+        }
+    }
 
     /**
      * Utility functions for TextClassifier methods.
diff --git a/core/java/android/view/textclassifier/TextClassifierImpl.java b/core/java/android/view/textclassifier/TextClassifierImpl.java
index d7aaee73f976..aea3cb0603f9 100644
--- a/core/java/android/view/textclassifier/TextClassifierImpl.java
+++ b/core/java/android/view/textclassifier/TextClassifierImpl.java
@@ -42,6 +42,9 @@ import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -66,6 +69,18 @@ final class TextClassifierImpl implements TextClassifier {
     private static final String MODEL_FILE_REGEX = "textclassifier\\.smartselection\\.(.*)\\.model";
     private static final String UPDATED_MODEL_FILE_PATH =
             "/data/misc/textclassifier/textclassifier.smartselection.model";
+    private static final List<String> ENTITY_TYPES_ALL =
+            Collections.unmodifiableList(Arrays.asList(
+                    TextClassifier.TYPE_ADDRESS,
+                    TextClassifier.TYPE_EMAIL,
+                    TextClassifier.TYPE_PHONE,
+                    TextClassifier.TYPE_URL));
+    private static final List<String> ENTITY_TYPES_BASE =
+            Collections.unmodifiableList(Arrays.asList(
+                    TextClassifier.TYPE_ADDRESS,
+                    TextClassifier.TYPE_EMAIL,
+                    TextClassifier.TYPE_PHONE,
+                    TextClassifier.TYPE_URL));
 
     private final Context mContext;
 
@@ -168,17 +183,23 @@ final class TextClassifierImpl implements TextClassifier {
 
     @Override
     public TextLinks generateLinks(
-            @NonNull CharSequence text, @NonNull TextLinks.Options options) {
+            @NonNull CharSequence text, @Nullable TextLinks.Options options) {
         Utils.validateInput(text);
         final String textString = text.toString();
         final TextLinks.Builder builder = new TextLinks.Builder(textString);
         try {
-            LocaleList defaultLocales = options != null ? options.getDefaultLocales() : null;
+            final LocaleList defaultLocales = options != null ? options.getDefaultLocales() : null;
+            final Collection<String> entitiesToIdentify =
+                    options != null && options.getEntityConfig() != null
+                            ? options.getEntityConfig().getEntities(this) : ENTITY_TYPES_ALL;
             final SmartSelection smartSelection = getSmartSelection(defaultLocales);
             final SmartSelection.AnnotatedSpan[] annotations = smartSelection.annotate(textString);
             for (SmartSelection.AnnotatedSpan span : annotations) {
-                final Map<String, Float> entityScores = new HashMap<>();
                 final SmartSelection.ClassificationResult[] results = span.getClassification();
+                if (results.length == 0 || !entitiesToIdentify.contains(results[0].mCollection)) {
+                    continue;
+                }
+                final Map<String, Float> entityScores = new HashMap<>();
                 for (int i = 0; i < results.length; i++) {
                     entityScores.put(results[i].mCollection, results[i].mScore);
                 }
@@ -193,6 +214,20 @@ final class TextClassifierImpl implements TextClassifier {
     }
 
     @Override
+    public Collection<String> getEntitiesForPreset(@TextClassifier.EntityPreset int entityPreset) {
+        switch (entityPreset) {
+            case TextClassifier.ENTITY_PRESET_NONE:
+                return Collections.emptyList();
+            case TextClassifier.ENTITY_PRESET_BASE:
+                return ENTITY_TYPES_BASE;
+            case TextClassifier.ENTITY_PRESET_ALL:
+                // fall through
+            default:
+                return ENTITY_TYPES_ALL;
+        }
+    }
+
+    @Override
     public void logEvent(String source, String event) {
         if (LOG_TAG.equals(source)) {
             mMetricsLogger.count(event, 1);
diff --git a/core/java/android/view/textclassifier/TextLinks.java b/core/java/android/view/textclassifier/TextLinks.java
index 4fe5662cca22..6c587cf9d475 100644
--- a/core/java/android/view/textclassifier/TextLinks.java
+++ b/core/java/android/view/textclassifier/TextLinks.java
@@ -161,11 +161,12 @@ public final class TextLinks {
     public static final class Options {
 
         private LocaleList mDefaultLocales;
+        private TextClassifier.EntityConfig mEntityConfig;
 
         /**
-         * @param defaultLocales ordered list of locale preferences that may be used to disambiguate
-         *      the provided text. If no locale preferences exist, set this to null or an empty
-         *      locale list.
+         * @param defaultLocales ordered list of locale preferences that may be used to
+         *                       disambiguate the provided text. If no locale preferences exist,
+         *                       set this to null or an empty locale list.
          */
         public Options setDefaultLocales(@Nullable LocaleList defaultLocales) {
             mDefaultLocales = defaultLocales;
@@ -173,6 +174,17 @@ public final class TextLinks {
         }
 
         /**
+         * Sets the entity configuration to use. This determines what types of entities the
+         * TextClassifier will look for.
+         *
+         * @param entityConfig EntityConfig to use
+         */
+        public Options setEntityConfig(@Nullable TextClassifier.EntityConfig entityConfig) {
+            mEntityConfig = entityConfig;
+            return this;
+        }
+
+        /**
          * @return ordered list of locale preferences that can be used to disambiguate
          *      the provided text.
          */
@@ -180,6 +192,15 @@ public final class TextLinks {
         public LocaleList getDefaultLocales() {
             return mDefaultLocales;
         }
+
+        /**
+         * @return The config representing the set of entities to look for.
+         * @see #setEntityConfig(TextClassifier.EntityConfig)
+         */
+        @Nullable
+        public TextClassifier.EntityConfig getEntityConfig() {
+            return mEntityConfig;
+        }
     }
 
     /**