Creating 'documents/lots/' endpoint

ereuse_devicehub/resources/documents/documents.py

@@ -2,6 +2,7 @@ import csv
 import datetime
 import enum
 import uuid
+from collections import OrderedDict
 from io import StringIO
 from typing import Callable, Iterable, Tuple
 
@@ -19,6 +20,8 @@ from ereuse_devicehub.resources.action import models as evs
 from ereuse_devicehub.resources.device import models as devs
 from ereuse_devicehub.resources.device.views import DeviceView
 from ereuse_devicehub.resources.documents.device_row import DeviceRow
+from ereuse_devicehub.resources.lot import LotView
+from ereuse_devicehub.resources.lot.models import Lot
 
 
 class Format(enum.Enum):
@@ -126,6 +129,41 @@ class DevicesDocumentView(DeviceView):
         return output
 
 
+class LotsDocumentView(LotView):
+    def find(self, args: dict):
+        return self.generate_lots_csv(Lot.query)
+
+    def generate_lots_csv(self, query):
+        """Get device query and put information in csv format."""
+        data = StringIO()
+        cw = csv.writer(data)
+        first = True
+        for lot in query:
+            l = LotRow(lot)
+            if first:
+                cw.writerow(l.keys())
+                first = False
+            cw.writerow(l.values())
+        output = make_response(data.getvalue())
+        output.headers['Content-Disposition'] = 'attachment; filename=lots-info.csv'
+        output.headers['Content-type'] = 'text/csv'
+        return output
+
+
+class LotRow(OrderedDict):
+    def __init__(self, lot: Lot) -> None:
+        super().__init__()
+        self.lot = lot
+        # General information about device
+        self['Id'] = lot.id.hex
+        self['Name'] = lot.name
+        self['Registered in'] = format(lot.created, '%c')
+        try:
+            self['Description'] = lot.description
+        except:
+            self['Description'] = ''
+
+
 class DocumentDef(Resource):
     __type__ = 'Document'
     SCHEMA = None
@@ -156,6 +194,8 @@ class DocumentDef(Resource):
         devices_view = DevicesDocumentView.as_view('devicesDocumentView',
                                                    definition=self,
                                                    auth=app.auth)
+        lots_view = LotsDocumentView.as_view('lotsDocumentView', definition=self)
         if self.AUTH:
             devices_view = app.auth.requires_auth(devices_view)
         self.add_url_rule('/devices/', defaults=d, view_func=devices_view, methods=get)
+        self.add_url_rule('/lots/', defaults=d, view_func=lots_view, methods=get)