Holger Frey
2 years ago
2 changed files with 31 additions and 1 deletions
@ -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…
Reference in new issue