Browse Source

added test for the `to_dict()` method

main
Holger Frey 2 years ago
parent
commit
c10ed0893b
  1. 2
      pytest.ini
  2. 30
      split_uniques.py

2
pytest.ini

@ -1,2 +1,2 @@ @@ -1,2 +1,2 @@
[pytest]
python_files = *.py
python_files = *.py

30
split_uniques.py

@ -0,0 +1,30 @@ @@ -0,0 +1,30 @@
import pandas as pd
from typing import Iterable, NamedTuple
SplitUniqueKeys = dict[str:str]
class SplitUniqueResult(NamedTuple):
keys: SplitUniqueKeys
data: pd.DataFrame
def split(
data: pd.DataFrame,
columns: str | Iterable[str],
*,
prevkeys: SplitUniqueKeys = None
) -> Iterable[SplitUniqueResult]:
if isinstance(columns, str):
columns = [columns]
if prevkeys is None:
prevkeys = {}
current, *rest = columns
for value in data[current].unique():
selection = data[current] == value
selected = data.loc[selection].copy()
keys = prevkeys | {current: value}
if rest:
yield from split(selected, rest, prevkeys=keys)
else:
yield SplitUniqueResult(keys, selected)
Loading…
Cancel
Save