Converting the write_output() function

In this function, we can see how to query a peewee model instance. On line 235, we need to select a count of files where the custodian is equal to the custodian's id. We first call select() on the model to signify we wish to select data, followed by the where() method to specify the column name, Files.custodian, and the value, custodian_model.id, to evaluate. This is followed by the count() method to provide an integer of the number of responsive results. Note that the count variable is an integer, not a tuple, like it was in the previous iteration:

226 def write_output(source, custodian_model):
227 """
228 The write_output function handles writing either the CSV or
229 HTML reports
230 :param source: The output filepath
231 :param custodian_model: Peewee model instance for the
232 custodian
233 :return: None
234 """
235 count = Files.select().where(
236 Files.custodian == custodian_model.id).count()
237
238 logger.info("{} files found for custodian.".format(count))

On line 240, we follow the same logic from the prior iteration to check and see whether some lines were responsive, followed by statements to validate the output extension to engage the correct writer or provide the user's accurate error information. Note that, this time, we pass along the custodian model instance versus an id or name on lines 243 and 247, as Peewee performs operations best on existing model instances:

240     if not count:
241 logger.error('Files not found for custodian')
242 elif source.endswith('.csv'):
243 write_csv(source, custodian_model)
244 elif source.endswith('.html'):
245 write_html(source, custodian_model)
246 elif not (source.endswith('.html') or \
247 source.endswith('.csv')):
248 logger.error('Could not determine file type')
249 else:
250 logger.error('Unknown Error Occurred')