fix(pm4py): fixed typing of get_variants methods

process-intelligence-solutions · Feb 22, 2024 · 64ce334 · 64ce334
1 parent d8e2a66
commit 64ce334
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pm4py/stats.py b/pm4py/stats.py
@@ -193,7 +193,7 @@ def get_trace_attribute_values(log: Union[EventLog, pd.DataFrame], attribute: st
         return ret
 
 
-def get_variants(log: Union[EventLog, pd.DataFrame], activity_key: str = "concept:name", timestamp_key: str = "time:timestamp", case_id_key: str = "case:concept:name") -> Dict[Tuple[str], List[Trace]]:
+def get_variants(log: Union[EventLog, pd.DataFrame], activity_key: str = "concept:name", timestamp_key: str = "time:timestamp", case_id_key: str = "case:concept:name") -> Union[Dict[Tuple[str], List[Trace]], Dict[Tuple[str], int]]:
     """
     Gets the variants from the log
 
@@ -212,7 +212,7 @@ def get_variants(log: Union[EventLog, pd.DataFrame], activity_key: str = "concep
     return get_variants_as_tuples(log, activity_key=activity_key, timestamp_key=timestamp_key, case_id_key=case_id_key)
 
 
-def get_variants_as_tuples(log: Union[EventLog, pd.DataFrame], activity_key: str = "concept:name", timestamp_key: str = "time:timestamp", case_id_key: str = "case:concept:name") -> Dict[Tuple[str], List[Trace]]:
+def get_variants_as_tuples(log: Union[EventLog, pd.DataFrame], activity_key: str = "concept:name", timestamp_key: str = "time:timestamp", case_id_key: str = "case:concept:name") -> Union[Dict[Tuple[str], List[Trace]], Dict[Tuple[str], int]]:
     """
     Gets the variants from the log (where the keys are tuples and not strings)