From e9dd37981d81249e3d0ea252eb7b3ecba4e71c2d Mon Sep 17 00:00:00 2001
From: Steve Lee <weichunstevelee@users.noreply.github.com>
Date: Tue, 11 Jul 2023 12:38:23 -0700
Subject: [PATCH] Expand Java definition extraction for TensorEncoderDef and
 InputSelectorDef (#71)

* Expand java extractor capability for InputSelectorDef and TensorEncoderDef

* Format file

---------

Co-authored-by: Steve Lee <stevel@netflix.com>
---
 .../buildgen/javadefref/JavaSourceEntityExtractor.scala   | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala b/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
index 2d9dc21b..f6420014 100644
--- a/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
+++ b/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
@@ -61,7 +61,13 @@ object JavaSourceEntityExtractor {
 
     // TODO there needs to be some way to configure this
     val generators: Set[String] =
-      Set("InputTransformDef", "FeatureEncoderDef", "DataTransformDef")
+      Set(
+        "InputTransformDef",
+        "FeatureEncoderDef",
+        "DataTransformDef",
+        "TensorEncoderDef",
+        "InputSelectorDef"
+      )
 
     val topLevelDefsTypes: SortedSet[Entity] =
       compUnit