refactor out some parent classes to remove duplication

hjwp · hjwp · commit 1223c1896dce · 2019-02-08T11:00:34.000Z
diff --git a/domain_model.py b/domain_model.py
@@ -2,15 +2,15 @@
 from typing import List
 from datetime import date
 
-
 @dataclass
 class OrderLine:
     sku: str
     quantity: int
 
 
+
 @dataclass
-class Order:
+class SkuLines:
     lines: List[OrderLine]
 
     @property
@@ -21,6 +21,12 @@ def skus(self):
     def quantities(self):
         return {l.sku: l.quantity for l in self.lines}
 
+
+
+
+@dataclass
+class Order(SkuLines):
+
     @property
     def fully_allocated(self):
         return self.allocation.is_complete
@@ -37,21 +43,9 @@ def allocate(self, stock, shipments):
         self.allocation.apply()
 
 
-class StockLine(OrderLine):
-    pass
-
 
 @dataclass
-class Stock:
-    lines: List[StockLine]
-
-    @property
-    def skus(self):
-        return set(l.sku for l in self.lines)
-
-    @property
-    def quantities(self):
-        return {l.sku: l.quantity for l in self.lines}
+class Stock(SkuLines):
 
     def can_allocate(self, line: OrderLine):
         return line.sku in self.skus and self.quantities[line.sku] > line.quantity
@@ -60,6 +54,8 @@ def allocate(self, sku, quantity):
         for line in self.lines:
             if line.sku == sku:
                 line.quantity -= quantity
+                return
+        raise Exception(f'sku {sku} not found in stock skus ({self.skus})')
 
 
 @dataclass
@@ -90,7 +86,7 @@ def sources(self):
         return {l.sku: l.source for l in self.lines}
 
     @staticmethod
-    def for_(order, source):
+    def for_(order: Order, source: Stock):
         return Allocation(
             lines=[
                 AllocationLine(sku=line.sku, source=source) for line in order.lines
@@ -113,4 +109,3 @@ def apply(self):
         for line in self.lines:
             line.source.allocate(line.sku, self.order.quantities[line.sku])
 
-
diff --git a/test_allocation.py b/test_allocation.py
@@ -6,9 +6,9 @@
 later = tomorrow + timedelta(days=10)
 
 Order = lambda d: domain_model.Order(lines=[domain_model.OrderLine(sku, qty) for sku, qty in d.items()])
-Stock = lambda d: domain_model.Stock(lines=[domain_model.StockLine(sku, qty) for sku, qty in d.items()])
+Stock = lambda d: domain_model.Stock(lines=[domain_model.OrderLine(sku, qty) for sku, qty in d.items()])
 Shipment = lambda d, eta: domain_model.Shipment(
-    lines=[domain_model.StockLine(sku, qty) for sku, qty in d.items()],
+    lines=[domain_model.OrderLine(sku, qty) for sku, qty in d.items()],
     eta=eta,
 )
 
