Coverage for utilities/export.py: 30%

1import csv 

2 

3from django.http import StreamingHttpResponse 

4from django.utils.encoding import force_str 

5from django.utils.http import content_disposition_header 

6from django.utils.translation import gettext_lazy as _ 

7from django_tables2.data import TableQuerysetData 

8from django_tables2.export import TableExport as TableExport_ 

9from django_tables2.rows import BoundRow 

10 

11from utilities.constants import CSV_DELIMITERS 

12 

13__all__ = ( 

14 'TableExport', 

15 'stream_table_csv_response', 

16) 

17 

18EXPORT_CHUNK_SIZE = 1000 

19 

20 

21class TableExport(TableExport_): 

22 """ 

23 A subclass of django-tables2's TableExport class which allows us to specify a delimiting 

24 characters for CSV exports. 

25 """ 

26 def __init__(self, *args, delimiter=None, **kwargs): 

27 if delimiter and delimiter not in CSV_DELIMITERS.keys(): 

28 raise ValueError(_("Invalid delimiter name: {name}").format(name=delimiter)) 

29 self.delimiter = delimiter or 'comma' 

30 super().__init__(*args, **kwargs) 

31 

32 def export(self): 

33 if self.format == self.CSV and self.delimiter is not None: 

34 delimiter = CSV_DELIMITERS[self.delimiter] 

35 return self.dataset.export(self.format, delimiter=delimiter) 

36 return super().export() 

37 

38 

39class _EchoBuffer: 

40 """ 

41 File-like object whose write() simply returns the value written, so csv.writer output can be 

42 captured row-by-row and fed to a StreamingHttpResponse. 

43 """ 

44 def write(self, value): 

45 return value 

46 

47 

48def stream_table_csv_response(table, exclude_columns=None, filename=None, delimiter=None): 

49 """ 

50 Return a StreamingHttpResponse that emits the given table's rows as CSV without first buffering 

51 the entire dataset in memory. Queryset-backed tables are iterated in chunks using 

52 QuerySet.iterator() to cap peak memory; prefetched relations are preserved under Django 4.1+. 

53 

54 Args: 

55 table: The django-tables2 Table instance to export 

56 exclude_columns: Iterable of column names to omit from the export 

57 filename: If set, a Content-Disposition header is included in the HTTP response, indicating its treatment 

58 as a file attachment with the specified name. 

59 delimiter: Name of a delimiter in utilities.constants.CSV_DELIMITERS (defaults to 'comma') 

60 """ 

61 if delimiter and delimiter not in CSV_DELIMITERS: 

62 raise ValueError(_("Invalid delimiter name: {name}").format(name=delimiter)) 

63 csv_delimiter = CSV_DELIMITERS[delimiter or 'comma'] 

64 exclude_columns = exclude_columns or set() 

65 

66 columns = [ 

67 column for column in table.columns.iterall() 

68 if not (column.column.exclude_from_export or column.name in exclude_columns) 

69 ] 

70 

71 writer = csv.writer(_EchoBuffer(), delimiter=csv_delimiter) 

72 

73 def iter_records(): 

74 if isinstance(table.data, TableQuerysetData): 

75 yield from table.data.data.iterator(chunk_size=EXPORT_CHUNK_SIZE) 

76 else: 

77 yield from table.data 

78 

79 def row_generator(): 

80 yield writer.writerow([ 

81 force_str(column.header, strings_only=True) for column in columns 

82 ]) 

83 for record in iter_records(): 

84 row = BoundRow(record, table=table) 

85 yield writer.writerow([ 

86 force_str(row.get_cell_value(column.name), strings_only=True) 

87 for column in columns 

88 ]) 

89 

90 response = StreamingHttpResponse(row_generator(), content_type='text/csv; charset=utf-8') 

91 if filename is not None: 

92 response['Content-Disposition'] = content_disposition_header(as_attachment=True, filename=filename) 

93 return response