Coverage for src/crimepy/combine.py: 97%

1"""Combines tables of crime data. 

2 

3`pandas.concat` turns a column of categories into plain objects when the 

4tables have different categories (as tables of different years usually do), 

5and `pandas.merge` does the same to the keys it joins on. These functions 

6keep the categories, so combined tables stay small and keep their labels. 

7 

8They take `amos` datasets or data frames and return an `amos.Dataset`, which 

9keeps the label, task, and groups of the first dataset given when its 

10columns are still there. 

11 

12Contents: 

13 merge: joins two tables, keeping their categories. 

14 stack: stacks tables on top of each other, keeping their categories. 

15 

16""" 

17 

18from __future__ import annotations 

19 

20from collections.abc import Iterable, Sequence 

21from typing import Any 

22 

23import amos 

24import pandas as pd 

25from pandas.api.types import union_categoricals 

26 

27from . import utilities 

28 

29""" Public Functions """ 

30 

31 

32def stack( 

33 tables: Iterable[amos.Dataset | pd.DataFrame], 

34 *, 

35 source: str | None = None, 

36 labels: Sequence[Any] | None = None) -> amos.Dataset: 

37 """Stacks tables on top of each other, keeping their categories. 

38 

39 Columns that are in some tables but not others are missing (`NA`) in the 

40 rows of the others. A column of categories in every table that has it 

41 stays a column of categories, with all of their categories. 

42 

43 Args: 

44 tables: the tables to stack (such as one table for each year). 

45 source: name of a column to add that says which table each row came 

46 from. Defaults to `None`, which adds no column. 

47 labels: what `source` says for each table. Defaults to `None`, in 

48 which case the tables are numbered from 0. 

49 

50 Raises: 

51 ValueError: if there are no tables, or `labels` does not have one 

52 label for each table. 

53 

54 Returns: 

55 A dataset of the stacked table, with a new index. 

56 

57 """ 

58 items = list(tables) 

59 frames = [utilities.to_frame(t) for t in items] 

60 if not frames: 

61 message = 'there are no tables to stack' 

62 raise ValueError(message) 

63 if labels is not None and len(labels) != len(frames): 

64 message = 'labels must have one label for each table' 

65 raise ValueError(message) 

66 if source is not None: 

67 names = list(labels) if labels is not None else list(range(len(frames))) 

68 frames = [ 

69 f.assign(**{source: name}) 

70 for f, name in zip(frames, names, strict = True)] 

71 columns = list(dict.fromkeys(c for f in frames for c in f.columns)) 

72 categorical = [ 

73 c for c in columns 

74 if all( 

75 isinstance(f[c].dtype, pd.CategoricalDtype) 

76 for f in frames if c in f.columns)] 

77 for column in categorical: 

78 union = union_categoricals( 

79 [f[column].array for f in frames if column in f.columns], 

80 ignore_order = True) 

81 frames = [ 

82 f.assign(**{column: f[column].cat.set_categories(union.categories)}) 

83 if column in f.columns else f 

84 for f in frames] 

85 stacked = pd.concat(frames, ignore_index = True) 

86 for column in categorical: 

87 if not isinstance(stacked[column].dtype, pd.CategoricalDtype): 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true

88 stacked[column] = stacked[column].astype('category') 

89 return _dataset(stacked, items) 

90 

91 

92def merge( 

93 left: amos.Dataset | pd.DataFrame, 

94 right: amos.Dataset | pd.DataFrame, 

95 on: Sequence[str] | str, 

96 *, 

97 how: str = 'left', 

98 suffixes: tuple[str, str] = ('', '_right'), 

99 validate: str | None = None) -> amos.Dataset: 

100 """Joins two tables, keeping the categories of their columns. 

101 

102 Args: 

103 left: the first table. 

104 right: the second table. 

105 on: the columns to join on, which both tables have. 

106 how: "left", "right", "inner", or "outer". Defaults to "left". 

107 suffixes: added to the names of columns (other than `on`) that both 

108 tables have: the first to those of `left`, the second to those of 

109 `right`. Defaults to ('', '_right'). 

110 validate: checks the kind of join, as in `pandas.merge` (such as 

111 "many_to_one"). Defaults to `None`, which checks nothing. 

112 

113 Returns: 

114 A dataset of the joined table. 

115 

116 """ 

117 keys = [on] if isinstance(on, str) else list(on) 

118 restore = {} 

119 items = [left, right] 

120 left = utilities.to_frame(left).copy() 

121 right = utilities.to_frame(right).copy() 

122 for key in keys: 

123 kinds = (left[key].dtype, right[key].dtype) 

124 if any(isinstance(k, pd.CategoricalDtype) for k in kinds): 

125 restore[key] = True 

126 left[key] = left[key].astype('object') 

127 right[key] = right[key].astype('object') 

128 joined = left.merge( 

129 right, on = keys, 

130 how = how, suffixes = suffixes, # type: ignore[arg-type] 

131 validate = validate) # type: ignore[arg-type] 

132 for key in restore: 

133 joined[key] = joined[key].astype('category') 

134 return _dataset(joined, items) 

135 

136 

137""" Private Functions """ 

138 

139 

140def _dataset( 

141 data: pd.DataFrame, 

142 items: Sequence[amos.Dataset | pd.DataFrame]) -> amos.Dataset: 

143 """Returns `data` as a dataset like the first dataset of `items`. 

144 

145 Args: 

146 data: the combined table. 

147 items: the datasets and data frames that were combined. 

148 

149 Returns: 

150 A dataset with the label, task, and groups of the first dataset in 

151 `items` whose label and groups are columns of `data` (and with 

152 none if no dataset was given). 

153 

154 """ 

155 first = next((i for i in items if isinstance(i, amos.Dataset)), None) 

156 if first is None: 

157 return amos.Dataset.create(data) 

158 label = first.label if first.label in data.columns else None 

159 groups = [g for g in first.groups if g in data.columns] 

160 return amos.Dataset.create( 

161 data, label = label, task = first.task if label else None, 

162 seed = first.seed, groups = groups or None)