refactoring helper classes into random-sets

pyproject.toml

@@ -14,6 +14,7 @@ typer = "^0.9.0"
 rich = "^13.7.0"
 pyyaml = "^6.0.1"
 csv2md = "^1.2.0"
+random_sets = { git = "file:///home/greg/dev/random-sets", branch="main" }
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.4.3"

rolltable/cli.py

@@ -1,4 +1,4 @@
-from rolltable import tables
+from rolltable.types import RollTable
 import typer
 from enum import Enum
 from rich import print
@@ -46,7 +46,7 @@ def create(
     CLI for creating roll tables.
     """
 
-    rt = tables.RollTable([Path(s).read_text() for s in sources], frequency=frequency, die=die, hide_rolls=hide_rolls)
+    rt = RollTable([Path(s).read_text() for s in sources], frequency=frequency, die=die, hide_rolls=hide_rolls)
 
     if output == OUTPUT_FORMATS.yaml:
         print(rt.as_yaml())

rolltable/types.py

@@ -1,151 +1,8 @@
 import yaml
-import random
 from csv2md.table import Table
 from collections.abc import Iterable
-from typing import Optional, List, IO, Union
+from typing import Optional, List, Union
-
+from random_sets.datasources import DataSource
-
-class DataSource:
-    """
-    Represents a yaml data source used to generate roll tables.
-
-    Attributes:
-
-        source      - the IO source to parse
-        frequency   - the frequency distribution to apply
-        headers     - an array of header strings
-        data        - The parsed YAML data
-
-    Methods:
-
-        load_source - Read and parse the source, populating the attributes
-
-    """
-    def __init__(self, source: IO, frequency: str = 'default') -> None:
-        """
-        Initialize a DataSource instance.
-
-        Args:
-            source      - an IO object to read source from
-            frequency   - the name of the frequency distribution to use; must
-                          be defined in the source file's metadata.
-        """
-        self.source = source
-        self.frequency = frequency
-        self.headers = []
-        self.frequencies = None
-        self.data = None
-        self.metadata = None
-        self.load_source()
-
-    def load_source(self) -> None:
-        """
-        Cache the yaml source and the parsed or generated metadata.
-        """
-        if self.data:
-            return
-        self.read_source()
-        self.init_headers()
-        self.init_frequencies()
-
-    def read_source(self) -> None:
-        self.data = yaml.safe_load(self.source)
-        self.metadata = self.data.pop('metadata', {})
-
-    def init_headers(self) -> None:
-        if 'headers' in self.metadata:
-            self.headers = self.metadata['headers']
-
-    def init_frequencies(self) -> None:
-        num_keys = len(self.data.keys())
-        default_freq = num_keys / 100
-
-        frequencies = {
-            'default': dict([(k, default_freq) for k in self.data.keys()])
-        }
-        if 'frequencies' in self.metadata:
-            frequencies.update(**self.metadata['frequencies'])
-        self.frequencies = frequencies[self.frequency]
-
-    def random_frequencies(self, count: int = 1) -> list:
-        """
-        Choose random option names from the frequency table.
-        """
-        weights = []
-        options = []
-        for (option, weight) in self.frequencies.items():
-            weights.append(weight)
-            options.append(option)
-        return random.choices(options, weights=weights, k=count)
-
-    def random_values(self, count: int = 1) -> list:
-        """
-        Return a list of random values from the data set, as a list of lists.
-        """
-        return [
-            self.get_entries(option, rand=True) for option in self.random_frequencies(count)
-        ]
-
-    def as_dict(self) -> dict:
-        """
-        Return the contents of the data source as a dict.
-        """
-        data = dict()
-        for name in self.data.keys():
-            entries = self.get_entries(name, rand=False)
-            items = {(k, v) for k, v in zip(self.headers, entries)}
-            data[name] = dict(items)
-        return data
-
-    def get_entries(self, option, rand: bool = False) -> list:
-        """
-        For a random item or each item in the specified option in the data source,
-        return a flattened list of the option, the select item, and the item's value (if any).
-        """
-
-        # If there is no data for the specified option, stop now.
-        flattened = [option]
-        if not self.data[option]:
-            return flattened
-
-        if hasattr(self.data[option], 'keys'):
-            # if the option is a dict, we assume the values are lists; we select a random item
-            # and prepend the key to the value list as our random selection. For example, given:
-            #
-            #  >>> self.data[option] == {'One': ['bar', 'baz'], 'Two': ['qaz', 'qux']}
-            #
-            # choices might then be: ['One', 'bar', 'baz']
-            #
-            if rand:
-                k, v = random.choice(list(self.data[option].items()))
-                choices = [[k] + v]
-            else:
-                choices = [
-                    [k] + v for k, v in list(self.data[option].items())
-                ]
-        else:
-            # If the option is either a list or a string, just select it.
-            choices = self.data[option]
-
-        for choice in choices:
-            # If the randomly-selected choice is a dict, choose a random item and return a list consisting
-            # of the option name, the key, and the value, flattening the # value if it is also a list.
-            if hasattr(choice, 'keys'):
-                for (k, v) in choice.items():
-                    if type(v) is list:
-                        flattened.extend([k, *v])
-                    else:
-                        flattened.extend([k, v])
-                continue
-
-            # if the member is a list, return the flattened list
-            if type(choice) is list:
-                flattened.extend(choice)
-                continue
-
-            # otherwise, return a list consisting of option and choice
-            flattened.append(choice)
-        return flattened
 
 
 class RollTable:

tests/test_tables.py

@@ -1,6 +1,6 @@
 import pytest
 
-from rolltable import tables
+from rolltable import types
 
 fixture_metadata = """
 metadata:
@@ -137,80 +137,80 @@ dict:
 
 @pytest.mark.parametrize('fixture', fixture_lists_and_dicts)
 def test_lists_and_dicts(fixture):
-    t = tables.RollTable([fixture], die=1)
+    t = types.RollTable([fixture], die=1)
     assert(str(t))
 
 
 def test_combined_tables():
-    combined = tables.RollTable([fixture_combined_A, fixture_combined_B], die=6)
+    combined = types.RollTable([fixture_combined_A, fixture_combined_B], die=6)
     assert str(combined)
 
 
 def test_table_end_to_end():
-    assert str(tables.RollTable([fixture_source]))
+    assert str(types.RollTable([fixture_source]))
 
 
 def test_table_end_to_end_with_metadata():
-    assert str(tables.RollTable([fixture_metadata + fixture_source]))
+    assert str(types.RollTable([fixture_metadata + fixture_source]))
 
 
 def test_table_frequency():
-    t = tables.RollTable([fixture_metadata + fixture_source], frequency='nondefault')
+    t = types.RollTable([fixture_metadata + fixture_source], frequency='nondefault')
     assert t._data[0].frequencies['Option 1'] == 0.0
     assert t._data[0].frequencies['Option 2'] == 0.1
     assert t._data[0].frequencies['Option 3'] == 0.9
 
 
 def test_one_option():
-    t = tables.RollTable([fixture_one_choice], die=1)
+    t = types.RollTable([fixture_one_choice], die=1)
     assert t._values == [['option 1', 'choice 1', 'description 1']]
 
 
 def test_collapsed():
-    t = tables.RollTable([fixture_repeated_choices], die=6)
+    t = types.RollTable([fixture_repeated_choices], die=6)
     assert len(list(t.rows)) == 2  # (+1 for headers)
 
 
 def test_not_collapsed():
-    t = tables.RollTable([fixture_repeated_choices], die=6)
+    t = types.RollTable([fixture_repeated_choices], die=6)
     assert len(list(t.expanded_rows)) == 7  # (+1 for headers)
 
 
 def test_no_descriptions():
-    t = tables.RollTable([fixture_no_descriptions], die=1)
+    t = types.RollTable([fixture_no_descriptions], die=1)
     assert 'd1' in str(t)
     assert 'option 1' in str(t)
 
 
 def test_no_options():
-    t = tables.RollTable([fixture_no_options])
+    t = types.RollTable([fixture_no_options])
     assert str(t)
 
 
 @pytest.mark.parametrize('table', [
-    tables.RollTable([fixture_no_options]),
+    types.RollTable([fixture_no_options]),
-    tables.RollTable([fixture_one_choice]),
+    types.RollTable([fixture_one_choice]),
-    tables.RollTable([fixture_metadata + fixture_source]),
+    types.RollTable([fixture_metadata + fixture_source]),
-    tables.RollTable([fixture_source]),
+    types.RollTable([fixture_source]),
 ])
 def test_yaml(table):
     assert table.as_yaml()
 
 
 def test_text():
-    assert repr(tables.RollTable([fixture_no_options]))
+    assert repr(types.RollTable([fixture_no_options]))
-    assert repr(tables.RollTable([fixture_one_choice]))
+    assert repr(types.RollTable([fixture_one_choice]))
-    assert repr(tables.RollTable([fixture_metadata + fixture_source]))
+    assert repr(types.RollTable([fixture_metadata + fixture_source]))
-    assert repr(tables.RollTable([fixture_source]))
+    assert repr(types.RollTable([fixture_source]))
 
 
 @pytest.mark.parametrize('table', [
-    tables.RollTable([fixture_no_options]),
+    types.RollTable([fixture_no_options]),
-    tables.RollTable([fixture_one_choice]),
+    types.RollTable([fixture_one_choice]),
-    tables.RollTable([fixture_metadata + fixture_source]),
+    types.RollTable([fixture_metadata + fixture_source]),
-    tables.RollTable([fixture_source]),
+    types.RollTable([fixture_source]),
-    tables.RollTable([fixture_no_options]),
+    types.RollTable([fixture_no_options]),
-    tables.RollTable([fixture_lists_and_dicts]),
+    types.RollTable([fixture_lists_and_dicts]),
 ])
 def test_as_dict(table):
     for src in table.datasources: