chore(datasets): Fix doctests (#488)

Signed-off-by: Merel Theisen <[email protected]>
kedro-org · Dec 19, 2023 · 6997b11 · 6997b11
1 parent f8f4a7d
commit 6997b11
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 14 deletions.
diff --git a/Makefile b/Makefile
@@ -37,17 +37,14 @@ dataset-doctest%:
 	# TODO(deepyaman): Fix as many doctests as possible (so that they run).
 	cd kedro-datasets && pytest kedro_datasets --doctest-modules --doctest-continue-on-failure --no-cov \
 	  --ignore kedro_datasets/databricks/managed_table_dataset.py \
-	  --ignore kedro_datasets/pandas/deltatable_dataset.py \
 	  --ignore kedro_datasets/pandas/gbq_dataset.py \
 	  --ignore kedro_datasets/partitions/incremental_dataset.py \
 	  --ignore kedro_datasets/partitions/partitioned_dataset.py \
-	  --ignore kedro_datasets/polars/lazy_polars_dataset.py \
 	  --ignore kedro_datasets/redis/redis_dataset.py \
 	  --ignore kedro_datasets/snowflake/snowpark_dataset.py \
 	  --ignore kedro_datasets/spark/deltatable_dataset.py \
 	  --ignore kedro_datasets/spark/spark_hive_dataset.py \
 	  --ignore kedro_datasets/spark/spark_jdbc_dataset.py \
-	  --ignore kedro_datasets/tensorflow/tensorflow_model_dataset.py \
 	  $(extra_pytest_arg${*})
 
 test-sequential:

diff --git a/kedro-datasets/kedro_datasets/databricks/managed_table_dataset.py b/kedro-datasets/kedro_datasets/databricks/managed_table_dataset.py
@@ -187,7 +187,7 @@ class ManagedTableDataset(AbstractVersionedDataset):
         ... )
         >>> data = [("Alex", 31), ("Bob", 12), ("Clarke", 65), ("Dave", 29)]
         >>> spark_df = SparkSession.builder.getOrCreate().createDataFrame(data, schema)
-        >>> dataset = ManagedTableDataset(table="names_and_ages")
+        >>> dataset = ManagedTableDataset(table="names_and_ages", write_mode="overwrite")
         >>> dataset.save(spark_df)
         >>> reloaded = dataset.load()
         >>> reloaded.take(4)

diff --git a/kedro-datasets/kedro_datasets/pandas/deltatable_dataset.py b/kedro-datasets/kedro_datasets/pandas/deltatable_dataset.py
@@ -73,7 +73,7 @@ class DeltaTableDataset(AbstractDataset):
         >>>
         >>> new_data = pd.DataFrame({"col1": [7, 8], "col2": [9, 10], "col3": [11, 12]})
         >>> dataset.save(new_data)
-        >>> dataset.get_loaded_version()
+        >>> assert isinstance(dataset.get_loaded_version(), int)
 
     """
 

diff --git a/kedro-datasets/kedro_datasets/pandas/gbq_dataset.py b/kedro-datasets/kedro_datasets/pandas/gbq_dataset.py
@@ -51,7 +51,7 @@ class GBQTableDataset(AbstractDataset[None, pd.DataFrame]):
         >>>
         >>> data = pd.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
         >>>
-        >>> dataset = GBQTableDataset("dataset", "table_name", project="my-project")
+        >>> dataset = GBQTableDataset(dataset="dataset", table_name="table_name", project="my-project")
         >>> dataset.save(data)
         >>> reloaded = dataset.load()
         >>>

diff --git a/kedro-datasets/kedro_datasets/polars/lazy_polars_dataset.py b/kedro-datasets/kedro_datasets/polars/lazy_polars_dataset.py
@@ -63,10 +63,10 @@ class LazyPolarsDataset(AbstractVersionedDataset[pl.LazyFrame, PolarsFrame]):
         >>>
         >>> data = pl.DataFrame({"col1": [1, 2], "col2": [4, 5], "col3": [5, 6]})
         >>>
-        >>> dataset = LazyPolarsDataset(filepath=tmp_path / "test.csv")
+        >>> dataset = LazyPolarsDataset(filepath=tmp_path / "test.csv", file_format="csv")
         >>> dataset.save(data)
         >>> reloaded = dataset.load()
-        >>> assert data.frame_equal(reloaded)
+        >>> assert data.frame_equal(reloaded.collect())
 
     """
 

diff --git a/kedro-datasets/kedro_datasets/spark/spark_jdbc_dataset.py b/kedro-datasets/kedro_datasets/spark/spark_jdbc_dataset.py
@@ -40,7 +40,7 @@ class SparkJDBCDataset(AbstractDataset[DataFrame, DataFrame]):
     .. code-block:: pycon
 
         >>> import pandas as pd
-        >>> from kedro_datasets.spark import SparkJBDCDataset
+        >>> from kedro_datasets.spark import SparkJDBCDataset
         >>> from pyspark.sql import SparkSession
         >>>
         >>> spark = SparkSession.builder.getOrCreate()

diff --git a/kedro-datasets/kedro_datasets/tensorflow/tensorflow_model_dataset.py b/kedro-datasets/kedro_datasets/tensorflow/tensorflow_model_dataset.py
@@ -50,14 +50,14 @@ class TensorFlowModelDataset(AbstractVersionedDataset[tf.keras.Model, tf.keras.M
         >>> import tensorflow as tf
         >>> import numpy as np
         >>>
-        >>> dataset = TensorFlowModelDataset(tmp_path / "data/06_models/tensorflow_model.h5")
-        >>> model = tf.keras.Model()
-        >>> predictions = model.predict([...])
+        >>> dataset = TensorFlowModelDataset(filepath=tmp_path / "data/06_models/tensorflow_model.h5")
+        >>> model = tf.keras.Sequential([tf.keras.layers.Dense(5, input_shape=(3,)),tf.keras.layers.Softmax()])
+        >>>
+        >>> # x = tf.random.uniform((10, 3))
+        >>> # predictions = model.predict(x)
         >>>
         >>> dataset.save(model)
         >>> loaded_model = dataset.load()
-        >>> new_predictions = loaded_model.predict([...])
-        >>> np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)
 
     """