Coverage for src/crimepy/cli.py: 84%

1"""The command line interface of crimepy. 

2 

3```sh 

4crimepy nibrs years 

5crimepy nibrs download 2021-2023 

6crimepy nibrs convert 2021-2023 

7crimepy nibrs export victims 2022 --states DE --output victims.csv 

8crimepy project study.ini --export 

9``` 

10 

11`export` writes a table of the master files (such as "agencies") or a built 

12table (such as "incidents"; see `crimepy.nibrs.tables`) to a csv or parquet 

13file, downloading and converting the years first if needed. `project` runs an 

14`amos` project from a settings file. 

15 

16Contents: 

17 main: runs the command line interface. 

18 

19""" 

20 

21from __future__ import annotations 

22 

23import argparse 

24import pathlib 

25import sys 

26from collections.abc import Sequence 

27 

28from . import nibrs 

29 

30""" Public Functions """ 

31 

32 

33def main(arguments: Sequence[str] | None = None) -> int: 

34 """Runs the command line interface. 

35 

36 Args: 

37 arguments: the arguments, without the program's name. Defaults to 

38 `None`, in which case they are those that the program was called 

39 with. 

40 

41 Returns: 

42 The exit status: 0 if the command succeeded. 

43 

44 """ 

45 parser = _parser() 

46 parsed = parser.parse_args(arguments) 

47 if not hasattr(parsed, 'run'): 

48 parser.print_help() 

49 return 1 

50 return int(parsed.run(parsed) or 0) 

51 

52 

53""" Private Functions """ 

54 

55 

56def _parser() -> argparse.ArgumentParser: 

57 """Returns the parser of the command line arguments. 

58 

59 Returns: 

60 The parser, with a subcommand for each command, each of which sets 

61 "run" to the function that runs it. 

62 

63 """ 

64 parser = argparse.ArgumentParser( 

65 prog = 'crimepy', 

66 description = ( 

67 'Download, combine, munge, analyze, and plot crime data.')) 

68 commands = parser.add_subparsers(title = 'commands') 

69 data = commands.add_parser('nibrs', help = 'work with NIBRS master files') 

70 tasks = data.add_subparsers(title = 'NIBRS commands') 

71 

72 years = tasks.add_parser('years', help = 'list the years the FBI has') 

73 years.set_defaults(run = _years) 

74 

75 download = tasks.add_parser('download', help = 'download master files') 

76 _add_years(download) 

77 _add_folder(download) 

78 download.add_argument( 

79 '--overwrite', action = 'store_true', 

80 help = 'download files that were downloaded before') 

81 download.set_defaults(run = _download) 

82 

83 convert = tasks.add_parser( 

84 'convert', help = 'convert master files to parquet files') 

85 _add_years(convert) 

86 _add_folder(convert) 

87 convert.add_argument( 

88 '--tables', nargs = '+', default = None, 

89 help = 'the tables to convert (all of them by default)') 

90 convert.add_argument( 

91 '--overwrite', action = 'store_true', 

92 help = 'convert tables that were converted before') 

93 convert.set_defaults(run = _convert) 

94 

95 export = tasks.add_parser( 

96 'export', help = 'write a table to a csv or parquet file') 

97 export.add_argument( 

98 'table', 

99 choices = sorted({*nibrs.layouts.SEGMENTS, *nibrs.tables.BUILDERS}), 

100 help = 'the table to write') 

101 _add_years(export) 

102 _add_folder(export) 

103 export.add_argument( 

104 '--states', nargs = '+', default = None, 

105 help = 'postal abbreviations of the states to keep') 

106 export.add_argument( 

107 '--codes', action = 'store_true', 

108 help = 'keep codes instead of replacing them with labels') 

109 export.add_argument( 

110 '--output', required = True, type = pathlib.Path, 

111 help = 'the file to write (.csv or .parquet)') 

112 export.set_defaults(run = _export) 

113 

114 project = commands.add_parser( 

115 'project', help = 'run an amos project from a settings file') 

116 project.add_argument('settings', type = pathlib.Path) 

117 project.add_argument( 

118 '--export', action = 'store_true', 

119 help = 'export the results to a folder named for the run') 

120 project.set_defaults(run = _project) 

121 return parser 

122 

123 

124def _add_folder(parser: argparse.ArgumentParser) -> None: 

125 """Adds the "--folder" argument to `parser`. 

126 

127 Args: 

128 parser: the parser of a command that uses the data folder. 

129 

130 """ 

131 parser.add_argument( 

132 '--folder', type = pathlib.Path, default = None, 

133 help = 'the data folder (by default, the "crime" folder of the data ' 

134 'folder of amos\'s workspace, beside the "code" folder)') 

135 

136 

137def _add_years(parser: argparse.ArgumentParser) -> None: 

138 """Adds the "years" argument to `parser`. 

139 

140 Args: 

141 parser: the parser of a command that takes years. 

142 

143 """ 

144 parser.add_argument( 

145 'years', nargs = '+', 

146 help = 'years or ranges of years, such as 2022 or 2015-2023') 

147 

148 

149def _convert(arguments: argparse.Namespace) -> int: 

150 """Converts master files to parquet files. 

151 

152 Args: 

153 arguments: the parsed arguments of the "convert" command. 

154 

155 Returns: 

156 The exit status, 0. 

157 

158 """ 

159 converted = nibrs.convert( 

160 arguments.years, arguments.folder, tables = arguments.tables, 

161 overwrite = arguments.overwrite) 

162 for year, paths in converted.items(): 

163 folder = next(iter(paths.values())).parent 

164 print(f'{year}: {len(paths)} tables in {folder}') 

165 return 0 

166 

167 

168def _download(arguments: argparse.Namespace) -> int: 

169 """Downloads master files. 

170 

171 Args: 

172 arguments: the parsed arguments of the "download" command. 

173 

174 Returns: 

175 The exit status, 0. 

176 

177 """ 

178 for path in nibrs.download( 

179 arguments.years, arguments.folder, overwrite = arguments.overwrite): 

180 print(path) 

181 return 0 

182 

183 

184def _export(arguments: argparse.Namespace) -> int: 

185 """Writes a table to a csv or parquet file. 

186 

187 Args: 

188 arguments: the parsed arguments of the "export" command. 

189 

190 Returns: 

191 The exit status: 0 if the table was written, or 2 if the output is 

192 not a csv or parquet file. 

193 

194 """ 

195 suffix = arguments.output.suffix.lower() 

196 if suffix not in ('.csv', '.parquet'): 

197 print('the output must be a .csv or .parquet file', file = sys.stderr) 

198 return 2 

199 decode = not arguments.codes 

200 if arguments.table in nibrs.tables.BUILDERS: 

201 data = nibrs.tables.build( 

202 arguments.table, arguments.years, states = arguments.states, 

203 folder = arguments.folder, decode = decode).data 

204 else: 

205 data = nibrs.load_frame( 

206 arguments.table, arguments.years, states = arguments.states, 

207 folder = arguments.folder, decode = decode) 

208 arguments.output.parent.mkdir(parents = True, exist_ok = True) 

209 if suffix == '.csv': 

210 data.to_csv(arguments.output, index = False) 

211 else: 

212 data.to_parquet(arguments.output, index = False) 

213 print(f'wrote {len(data):,} rows to {arguments.output}') 

214 return 0 

215 

216 

217def _project(arguments: argparse.Namespace) -> int: 

218 """Runs an amos project from a settings file. 

219 

220 Args: 

221 arguments: the parsed arguments of the "project" command. 

222 

223 Returns: 

224 The exit status, 0. 

225 

226 """ 

227 import amos # noqa: PLC0415 

228 

229 project = amos.Project.create(arguments.settings) 

230 if arguments.export: 

231 print(project.export()) 

232 print(project.result) 

233 return 0 

234 

235 

236def _years(arguments: argparse.Namespace) -> int: # noqa: ARG001 

237 """Lists the years of NIBRS master files that the FBI has. 

238 

239 Args: 

240 arguments: the parsed arguments of the "years" command, which are 

241 not used. 

242 

243 Returns: 

244 The exit status, 0. 

245 

246 """ 

247 found = nibrs.available_years() 

248 print(f'{found[0]}-{found[-1]}') 

249 return 0