Feedback

Author: fealho

sdv/datasets/demo.py

@@ -386,16 +386,11 @@ def _get_text_file_content(modality, dataset_name, filename, output_filepath=Non
     _validate_modalities(modality)
     if output_filepath is not None and not str(output_filepath).endswith('.txt'):
         fname = (filename or '').lower()
-        if 'readme' in fname:
-            raise ValueError(
-                'The README can only be saved as a txt file. '
-                "Please provide a filepath ending in '.txt'"
-            )
-        if 'source' in fname:
-            raise ValueError(
-                'The source can only be saved as a txt file. '
-                "Please provide a filepath ending in '.txt'"
-            )
+        file_type = 'README' if 'readme' in fname else 'source'
+        raise ValueError(
+            f'The {file_type} can only be saved as a txt file. '
+            "Please provide a filepath ending in '.txt'"
+        )
 
     dataset_prefix = f'{modality}/{dataset_name}/'
     contents = _list_objects(dataset_prefix)

tests/unit/datasets/test_demo.py

@@ -762,11 +762,11 @@ def fake(modality, dataset_name, filename, output_filepath=None):
     monkeypatch.setattr('sdv.datasets.demo._get_text_file_content', fake)
 
     # Run
-    r = get_readme('single_table', 'dataset1', '/tmp/readme.txt')
-    s = get_source('single_table', 'dataset1', '/tmp/source.txt')
+    readme = get_readme('single_table', 'dataset1', '/tmp/readme.txt')
+    source = get_source('single_table', 'dataset1', '/tmp/source.txt')
 
     # Assert
-    assert r == 'X' and s == 'X'
+    assert readme == 'X' and source == 'X'
     assert calls[0] == ('single_table', 'dataset1', 'README.txt', '/tmp/readme.txt')
     assert calls[1] == ('single_table', 'dataset1', 'SOURCE.txt', '/tmp/source.txt')