From b2a7e0ba66c034d6181b33ac061186bc05ed0fe6 Mon Sep 17 00:00:00 2001
From: Fabio Buso <fabio@logicalclocks.com>
Date: Wed, 13 Jan 2021 13:18:58 +0100
Subject: [PATCH] Add Python training dataset API reference (#215)

---
 auto_doc.py                                | 11 +++++++++++
 docs/templates/api/training_dataset_api.md | 19 +++++++++++++++++++
 mkdocs.yml                                 |  1 +
 3 files changed, 31 insertions(+)
 create mode 100644 docs/templates/api/training_dataset_api.md

diff --git a/auto_doc.py b/auto_doc.py
index 47e51b42ab..c1259631b5 100644
--- a/auto_doc.py
+++ b/auto_doc.py
@@ -102,6 +102,17 @@
         ),
         "fg_methods": keras_autodoc.get_methods("hsfs.feature_group.FeatureGroup"),
     },
+    "api/training_dataset_api.md": {
+        "td": ["hsfs.training_dataset.TrainingDataset"],
+        "td_create": ["hsfs.feature_store.FeatureStore.create_training_dataset"],
+        "td_get": ["hsfs.feature_store.FeatureStore.get_training_dataset"],
+        "td_properties": keras_autodoc.get_properties(
+            "hsfs.training_dataset.TrainingDataset"
+        ),
+        "td_methods": keras_autodoc.get_methods(
+            "hsfs.training_dataset.TrainingDataset"
+        ),
+    },
     "api/feature_api.md": {
         "feature": ["hsfs.feature.Feature"],
         "feature_properties": keras_autodoc.get_properties("hsfs.feature.Feature"),
diff --git a/docs/templates/api/training_dataset_api.md b/docs/templates/api/training_dataset_api.md
new file mode 100644
index 0000000000..a536964657
--- /dev/null
+++ b/docs/templates/api/training_dataset_api.md
@@ -0,0 +1,19 @@
+# Training Dataset
+
+{{td}}
+
+## Creation
+
+{{td_create}}
+
+## Retrieval
+
+{{td_get}}
+
+## Properties
+
+{{td_properties}}
+
+## Methods
+
+{{td_methods}}
diff --git a/mkdocs.yml b/mkdocs.yml
index 838c5ee322..03a1b953d1 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -57,6 +57,7 @@ nav:
     - Connection: generated/api/connection_api.md
     - FeatureStore: generated/api/feature_store_api.md
     - FeatureGroup: generated/api/feature_group_api.md
+    - TrainingDataset: generated/api/training_dataset_api.md
     - Storage Connector: generated/api/storage_connector_api.md
     - Feature: generated/api/feature_api.md
   - Contributing: CONTRIBUTING.md