Coverage for src/backend/InvenTree/plugin/base/integration/DataExport.py: 47%

1"""Plugin class for custom data exporting.""" 

2 

3from collections import OrderedDict 

4from typing import Optional 

5 

6from django.contrib.auth.models import User 

7from django.db.models import QuerySet 

8 

9from rest_framework import serializers, views 

10 

11from common.models import DataOutput 

12from InvenTree.helpers import current_date 

13from plugin import PluginMixinEnum 

14 

15 

16class DataExportMixin: 

17 """Mixin which provides ability to customize data exports. 

18 

19 When exporting data from the API, this mixin can be used to provide 

20 custom data export functionality. 

21 """ 

22 

23 ExportOptionsSerializer = None 

24 

25 # How many rows to fetch into memory at once when exporting data? 

26 EXPORT_CHUNK_SIZE: int = 250 

27 

28 class MixinMeta: 

29 """Meta options for this mixin.""" 

30 

31 MIXIN_NAME = 'DataExport' 

32 

33 def __init__(self): 

34 """Register mixin.""" 

35 super().__init__() 

36 self.add_mixin(PluginMixinEnum.EXPORTER, True, __class__) 

37 

38 def supports_export( 

39 self, 

40 model_class: type, 

41 user: User, 

42 serializer_class: Optional[serializers.Serializer] = None, 

43 view_class: Optional[views.APIView] = None, 

44 *args, 

45 **kwargs, 

46 ) -> bool: 

47 """Return True if this plugin supports exporting data for the given model. 

48 

49 Args: 

50 model_class: The model class to check 

51 user: The user requesting the export 

52 serializer_class: The serializer class to use for exporting the data 

53 view_class: The view class to use for exporting the data 

54 

55 Returns: 

56 True if the plugin supports exporting data for the given model 

57 """ 

58 # By default, plugins support all models 

59 return True 

60 

61 def generate_filename(self, model_class, export_format: str) -> str: 

62 """Generate a filename for the exported data.""" 

63 model = model_class.__name__ 

64 date = current_date().isoformat() 

65 

66 return f'InvenTree_{model}_{date}.{export_format}' 

67 

68 def update_headers( 

69 self, headers: OrderedDict, context: dict, **kwargs 

70 ) -> OrderedDict: 

71 """Update the headers for the data export. 

72 

73 Allows for optional modification of the headers for the data export. 

74 

75 Arguments: 

76 headers: The current headers for the export 

77 context: The context for the export (provided by the plugin serializer) 

78 

79 Returns: The updated headers 

80 """ 

81 # The default implementation does nothing 

82 return headers 

83 

84 def filter_queryset(self, queryset: QuerySet) -> QuerySet: 

85 """Filter the queryset before exporting data.""" 

86 # The default implementation returns the queryset unchanged 

87 return queryset 

88 

89 def export_data( 

90 self, 

91 queryset: QuerySet, 

92 serializer_class: serializers.Serializer, 

93 headers: OrderedDict, 

94 context: dict, 

95 output: DataOutput, 

96 serializer_context: Optional[dict] = None, 

97 **kwargs, 

98 ) -> list: 

99 """Export data from the queryset. 

100 

101 This method should be implemented by the plugin to provide 

102 the actual data export functionality. 

103 

104 Arguments: 

105 queryset: The queryset to export 

106 serializer_class: The serializer class to use for exporting the data 

107 serializer_context: Optional context for the serializer 

108 headers: The headers for the export 

109 context: Any custom context for the export (provided by the plugin serializer) 

110 output: The DataOutput object for the export 

111 

112 Returns: The exported data (a list of dict objects) 

113 """ 

114 output.refresh_from_db() 

115 N = queryset.count() 

116 

117 rows = [] 

118 

119 offset = 0 

120 

121 while offset < N: 

122 chunk = queryset[offset : offset + self.EXPORT_CHUNK_SIZE] 

123 

124 chunk_rows = serializer_class( 

125 chunk, many=True, exporting=True, context=serializer_context or {} 

126 ).data 

127 

128 rows.extend(chunk_rows) 

129 

130 offset += self.EXPORT_CHUNK_SIZE 

131 

132 # Update the export progress 

133 output.progress += len(chunk_rows) 

134 output.save() 

135 

136 return rows 

137 

138 def get_export_options_serializer(self, **kwargs) -> serializers.Serializer | None: 

139 """Return a serializer class with dynamic export options for this plugin. 

140 

141 Returns: 

142 A class instance of a DRF serializer class, by default this an instance of 

143 self.ExportOptionsSerializer using the *args, **kwargs if existing for this plugin 

144 """ 

145 # By default, look for a class level attribute 

146 serializer = getattr(self, 'ExportOptionsSerializer', None) 

147 

148 if serializer: 

149 return serializer(**kwargs)