bundle models have arrived

app/helpers/service_helper.py

@@ -2,7 +2,7 @@ import csv
 import io
 import re
 
-def items_csv_to_dict(items_csv_reader, irt_model):
+def items_csv_to_dict(items_csv_reader):
   items = []
   headers = []
 

app/helpers/solver_helper.py

@@ -25,12 +25,12 @@ def build_constraints(solver_run, problem, items):
                         * items[item.id]
                         for item in solver_run.items]) <= round(total_form_items * (max / 100)), f'{attribute.id} - {attribute.value} - max'
       elif attribute.type == 'bundle':
-        bundles = solver_run.bundles(attribute.id)
+        # TODO: account for many different bundle types, since the id condition in L33 could yield duplicates
         total_bundles = randint(constraint.minimum, constraint.maximum)
-        selected_bundles = sample(bundles, total_bundles)
+        selected_bundles = sample(solver_run.bundles, total_bundles)
 
         for bundle in selected_bundles:
-          problem += lpSum([items[item.id] for item in solver_run.items if getattr(item, bundle["type"], None) == bundle['id']]) == bundle['count'], f'Bundle constraint for {bundle["type"]} ({bundle["id"]})'
+          problem += lpSum([items[item.id] for item in solver_run.items if getattr(item, bundle.type, None) == bundle.id]) == bundle.count, f'Bundle constraint for {bundle.type} ({bundle.id})'
 
 
     return problem

app/models/bundle.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+class Bundle(BaseModel):
+  id: int
+  count: int
+  type: str

app/models/solver_run.py

@@ -4,11 +4,13 @@ from typing import List, Optional
 from models.item import Item
 from models.constraint import Constraint
 from models.irt_model import IRTModel
+from models.bundle import Bundle
 from models.objective_function import ObjectiveFunction
 from models.advanced_options import AdvancedOptions
 
 class SolverRun(BaseModel):
   items: List[Item]
+  bundles: Optional[Bundle]
   constraints: List[Constraint]
   irt_model: IRTModel
   objective_function: ObjectiveFunction
@@ -28,25 +30,36 @@ class SolverRun(BaseModel):
     self.items = [item for item in self.items if item not in items]
     return True
 
-  def bundles(self, type_attribute):
-    bundles = []
+  def generate_bundles(self):
+    bundle_constraints = (constraint.reference_attribute for constraint in self.constraints if constraint.reference_attribute.type == 'bundle')
 
-    for item in self.items:
-      # dynamically get the attribute that will be used as an identifier to bundle like items
-      attribute_id = getattr(item, type_attribute, None)
+    for bundle_constraint in bundle_constraints:
+      type_attribute = bundle_constraint.id
 
-      # make sure the item has said attribute
-      if attribute_id != None:
-        # get index of the bundle in the bundles list
-        bundle_index = next((index for (index, bundle) in enumerate(bundles) if bundle['id'] == attribute_id), None)
+      for item in self.items:
+        attribute_id = getattr(item, type_attribute, None)
 
-        # if the bundle index isn't found then the bundle hasn't been created
-        # and added to the list and needs to be, else increment the count of items
-        # in the bundle
-        if bundle_index == None:
-          # TODO: create actual "bundle" models instead of dicts
-          bundles.append({ 'id': attribute_id, 'count': 1, 'type': type_attribute })
-        else:
-          bundles[bundle_index]['count'] += 1
+        # make sure the item has said attribute
+        if attribute_id != None:
+          # if there are pre-existing bundles, add new or increment existing
+          # else create array with new bundle
+          if self.bundles != None:
+            # get index of the bundle in the bundles list if exists or None if it doesn't
+            bundle_index = next((index for (index, bundle) in enumerate(self.bundles) if bundle.id == attribute_id and bundle.type == type_attribute), None)
 
-    return bundles
+            # if the index doesn't exist add the new bundle of whatever type
+            # else increment the count of the current bundle
+            if bundle_index == None:
+              self.bundles.append(Bundle(
+                id=attribute_id,
+                count=1,
+                type=type_attribute
+              ))
+            else:
+              self.bundles[bundle_index].count += 1
+          else:
+            self.bundles = [Bundle(
+              id=attribute_id,
+              count=1,
+              type=type_attribute
+            )]

app/services/loft_service.py

@@ -15,6 +15,7 @@ class LoftService(Base):
   def process(self):
     try:
       self.solver_run = SolverRun.parse_obj(self.retreive_attributes_from_message())
+      self.solver_run.generate_bundles()
       self.solution = self.generate_solution()
       self.result = self.stream_to_s3_bucket()
     except ItemGenerationError as error:
@@ -40,7 +41,7 @@ class LoftService(Base):
     items_csv_reader = csv_helper.file_stream_reader(items_csv)
 
     # add items to attributes dict
-    attributes['items'] = service_helper.items_csv_to_dict(items_csv_reader, attributes['irt_model'])
+    attributes['items'] = service_helper.items_csv_to_dict(items_csv_reader)
     logging.info('Processed Attributes...')
 
     return attributes