Collect the page info and builder classes into the root module.

PiperOrigin-RevId: 423357060

Author: MarkDaoust

tools/tensorflow_docs/api_generator/pretty_docs/__init__.py

@@ -13,3 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
+from tensorflow_docs.api_generator.pretty_docs.base_page import PageBuilder
+from tensorflow_docs.api_generator.pretty_docs.base_page import PageInfo
+from tensorflow_docs.api_generator.pretty_docs.base_page import TemplatePageBuilder
+
+from tensorflow_docs.api_generator.pretty_docs.class_page import ClassPageBuilder
+from tensorflow_docs.api_generator.pretty_docs.class_page import ClassPageInfo
+
+from tensorflow_docs.api_generator.pretty_docs.function_page import FunctionPageBuilder
+from tensorflow_docs.api_generator.pretty_docs.function_page import FunctionPageInfo
+
+from tensorflow_docs.api_generator.pretty_docs.module_page import ModulePageBuilder
+from tensorflow_docs.api_generator.pretty_docs.module_page import ModulePageInfo
+
+from tensorflow_docs.api_generator.pretty_docs.type_alias_page import TypeAliasPageBuilder
+from tensorflow_docs.api_generator.pretty_docs.type_alias_page import TypeAliasPageInfo

tools/tensorflow_docs/api_generator/pretty_docs/base_page.py

@@ -25,11 +25,6 @@
 
 import jinja2
 
-JINJA_ENV = jinja2.Environment(
-    trim_blocks=True,
-    lstrip_blocks=True,
-    loader=jinja2.FileSystemLoader(str(pathlib.Path(__file__).parent)))
-
 
 class PageBuilder(abc.ABC):
 
@@ -42,10 +37,17 @@ def build(self) -> str:
 
 
 class TemplatePageBuilder(PageBuilder):
+  """A Page builder implemented on a jinja template."""
+
   TEMPLATE = 'templates/page.jinja'
+  TEMPLATE_SEARCH_PATH = tuple([str(pathlib.Path(__file__).parent)])
+  JINJA_ENV = jinja2.Environment(
+      trim_blocks=True,
+      lstrip_blocks=True,
+      loader=jinja2.FileSystemLoader(TEMPLATE_SEARCH_PATH))
 
   def build(self) -> str:
-    template = JINJA_ENV.get_template(self.TEMPLATE)
+    template = self.JINJA_ENV.get_template(self.TEMPLATE)
     content = template.render(builder=self, page_info=self.page_info)
     return content