diff --git a/random_sets/datasources.py b/random_sets/datasources.py index 442e82f..c20c1ed 100644 --- a/random_sets/datasources.py +++ b/random_sets/datasources.py @@ -124,6 +124,7 @@ class DataSource: # If there is no data for the specified option, stop now. flattened = [] if not self.data[option]: + return [option] raise ValueError(f"There is no data for '{option}' in your data source.") if hasattr(self.data[option], 'keys'): diff --git a/test/test_datasources.py b/test/test_datasources.py index a6ec357..f0f2c43 100644 --- a/test/test_datasources.py +++ b/test/test_datasources.py @@ -56,6 +56,19 @@ def test_zero_frequency(): assert 'Option 1' not in val +def test_one_column_only(): + fixture = StringIO(f""" +{fixture_metadata} +Option 1: +Option 2: +Option 3: +""") + ds = datasources.DataSource(fixture) + vals = ds.random_values(count=10) + assert len(vals) == 10 + assert len(vals[0]) == 1 + + def test_distribution_accuracy_to_one_decimal_place(): fixture = StringIO(fixture_metadata + fixture_source) ds = datasources.DataSource(fixture)