wrap sum() into fn flatten_list

pyhealth/datasets/base_dataset.py

@@ -10,7 +10,8 @@
 
 from pyhealth.data import Patient, Event
 from pyhealth.datasets.utils import MODULE_CACHE_PATH
-from pyhealth.datasets.utils import hash_str, list_nested_level, is_homo_list
+from pyhealth.datasets.utils import hash_str
+from pyhealth.datasets.utils import list_nested_level, is_homo_list, flatten_list
 from pyhealth.medcode import CrossMap
 from pyhealth.utils import load_pickle, save_pickle
 
@@ -356,9 +357,7 @@ def set_task(
  # 1 level nested list
  if level == 1:
  # a list of values of the same type
- # sum() flattens the nested list
- # e.g, [[1, 2], [3, 4]] -> [1, 2, 3, 4]
- check = is_homo_list(sum([s[key] for s in samples], []))
+ check = is_homo_list(flatten_list([s[key] for s in samples]))
  assert check, \
  f"Key {key} has mixed types in the nested list within samples"
  # 2 level nested list
@@ -370,8 +369,9 @@ def set_task(
  assert all(check), \
  f"Key {key} has mixed nested list levels within samples"
  # a list of list of values of the same type
- # sum() flattens the nested list
- check = is_homo_list(sum([l for s in samples for l in s[key]], []))
+ check = is_homo_list(
+ flatten_list([l for s in samples for l in s[key]])
+ )
  assert check, \
  f"Key {key} has mixed types in the nested list within samples"
 
@@ -459,8 +459,7 @@ def get_all_tokens(
  continue
  # a list of lists of values
  elif type(sample[key][0]) == list:
- # sum() flattens the nested list
- tokens.extend(sum(sample[key], []))
+ tokens.extend(flatten_list(sample[key]))
  # a list of values
  else:
  tokens.extend(sample[key])
@@ -568,8 +567,9 @@ def task_stat(self) -> str:
  nested = [isinstance(e, list) for s in self.samples for e in s[key]]
  # key's feature is a list of lists
  if any(nested):
- # sum() flattens the nested list
- num_events = [len(sum(sample[key], [])) for sample in self.samples]
+ num_events = [
+ len(flatten_list(sample[key])) for sample in self.samples
+ ]
  # key's feature is a list of values
  else:
  num_events = [len(sample[key]) for sample in self.samples]

pyhealth/datasets/utils.py

@@ -33,6 +33,25 @@ def strptime(s: str) -> Optional[datetime]:
  return dateutil_parse(s)
 
 
+def flatten_list(l: List) -> List:
+ """Flattens a list of list.
+
+ Args:
+ l: List, the list of list to be flattened.
+
+ Returns:
+ List, the flattened list.
+
+ Examples:
+ >>> flatten_list([[1], [2, 3], [4]])
+ [1, 2, 3, 4]R
+ >>> flatten_list([[1], [[2], 3], [4]])
+ [1, [2], 3, 4]
+ """
+ assert isinstance(l, list), "l must be a list."
+ return sum(l, start=[])
+
+
 def list_nested_level(l: List) -> int:
  """Gets the nested level of a list.