Skip to content

Commit d9446f7

Browse files
hjwphjwp
committed
move away from dicts
1 parent 725ce5c commit d9446f7

2 files changed

Lines changed: 130 additions & 80 deletions

File tree

‎domain_model.py‎

Lines changed: 81 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,32 @@
1-
class Allocation(dict):
1+
from dataclasses import dataclass
2+
from typing import List
3+
from datetime import date
24

3-
def __init__(self, d, order):
4-
self.order = order
5-
super().__init__(d)
65

7-
@property
8-
def skus(self):
9-
return self.keys()
6+
@dataclass
7+
class OrderLine:
8+
sku: str
9+
quantity: int
1010

11-
@staticmethod
12-
def for_(order, source):
13-
return Allocation({
14-
sku: source
15-
for sku, quantity in order.items()
16-
if source.can_allocate(sku, quantity)
17-
}, order=order)
1811

19-
def supplement_with(self, allocation):
20-
for sku, quantity in allocation.items():
21-
if sku in self:
22-
continue
23-
self[sku] = quantity
12+
@dataclass
13+
class Order:
14+
lines: List[OrderLine]
2415

2516
@property
26-
def is_complete(self):
27-
return self.skus == self.order.skus
28-
29-
def apply(self):
30-
for sku, source in self.items():
31-
source[sku] -= self.order[sku]
32-
33-
34-
class Order(dict):
17+
def skus(self):
18+
return set(l.sku for l in self.lines)
3519

3620
@property
37-
def skus(self):
38-
return self.keys()
21+
def quantities(self):
22+
return {l.sku: l.quantity for l in self.lines}
3923

4024
@property
4125
def fully_allocated(self):
4226
return self.allocation.is_complete
4327

4428
def allocate(self, stock, shipments):
45-
self.allocation = Allocation({}, order=self)
29+
self.allocation = Allocation(lines=[], order=self)
4630
for source in [stock] + sorted(shipments):
4731
source_allocation = Allocation.for_(self, source)
4832
if source_allocation.is_complete:
@@ -53,21 +37,80 @@ def allocate(self, stock, shipments):
5337
self.allocation.apply()
5438

5539

56-
class Stock(dict):
40+
class StockLine(OrderLine):
41+
pass
42+
43+
44+
@dataclass
45+
class Stock:
46+
lines: List[StockLine]
47+
48+
@property
49+
def skus(self):
50+
return set(l.sku for l in self.lines)
51+
52+
@property
53+
def quantities(self):
54+
return {l.sku: l.quantity for l in self.lines}
5755

58-
def can_allocate(self, sku, quantity):
59-
return sku in self and self[sku] > quantity
56+
def can_allocate(self, line: OrderLine):
57+
return line.sku in self.skus and self.quantities[line.sku] > line.quantity
6058

6159
def allocate(self, sku, quantity):
62-
self[sku] -= quantity
60+
for line in self.lines:
61+
if line.sku == sku:
62+
line.quantity -= quantity
6363

6464

65+
@dataclass
6566
class Shipment(Stock):
67+
eta: date = None
6668

6769
def __lt__(self, other):
6870
return self.eta < other.eta
6971

70-
def __init__(self, lines, eta):
71-
self.eta = eta
72-
super().__init__(lines)
72+
73+
@dataclass
74+
class AllocationLine:
75+
sku: str
76+
source: Stock
77+
78+
79+
@dataclass
80+
class Allocation:
81+
lines: List[AllocationLine]
82+
order: Order
83+
84+
@property
85+
def skus(self):
86+
return set(l.sku for l in self.lines)
87+
88+
@property
89+
def sources(self):
90+
return {l.sku: l.source for l in self.lines}
91+
92+
@staticmethod
93+
def for_(order, source):
94+
return Allocation(
95+
lines=[
96+
AllocationLine(sku=line.sku, source=source) for line in order.lines
97+
if source.can_allocate(line)
98+
],
99+
order=order
100+
)
101+
102+
def supplement_with(self, allocation):
103+
for line in allocation.lines:
104+
if line.sku in self.skus:
105+
continue
106+
self.lines.append(line)
107+
108+
@property
109+
def is_complete(self):
110+
return self.skus == self.order.skus
111+
112+
def apply(self):
113+
for line in self.lines:
114+
line.source.allocate(line.sku, self.order.quantities[line.sku])
115+
73116

‎test_allocation.py‎

Lines changed: 49 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,26 @@
1-
from domain_model import Order, Stock, Shipment
1+
import domain_model
22
from datetime import date, timedelta
33

44
today = date.today()
55
tomorrow = today + timedelta(days=1)
66
later = tomorrow + timedelta(days=10)
77

8+
Order = lambda d: domain_model.Order(lines=[domain_model.OrderLine(sku, qty) for sku, qty in d.items()])
9+
Stock = lambda d: domain_model.Stock(lines=[domain_model.StockLine(sku, qty) for sku, qty in d.items()])
10+
Shipment = lambda d, eta: domain_model.Shipment(
11+
lines=[domain_model.StockLine(sku, qty) for sku, qty in d.items()],
12+
eta=eta,
13+
)
14+
815

916
def test_can_allocate_to_stock():
1017
order = Order({'a-sku': 10})
1118
stock = Stock({'a-sku': 1000})
1219

1320
order.allocate(stock, shipments=[])
1421

15-
assert order.allocation['a-sku'] == stock
16-
assert stock['a-sku'] == 990
22+
assert order.allocation.sources['a-sku'] == stock
23+
assert stock.quantities['a-sku'] == 990
1724

1825

1926
def test_can_allocate_to_shipment():
@@ -22,8 +29,8 @@ def test_can_allocate_to_shipment():
2229

2330
order.allocate(stock=Stock({}), shipments=[shipment])
2431

25-
assert order.allocation['a-sku'] == shipment
26-
assert shipment['a-sku'] == 990
32+
assert order.allocation.sources['a-sku'] == shipment
33+
assert shipment.quantities['a-sku'] == 990
2734

2835

2936
def test_ignores_irrelevant_stock():
@@ -33,7 +40,7 @@ def test_ignores_irrelevant_stock():
3340

3441
order.allocate(stock=stock, shipments=[shipment])
3542

36-
assert order.allocation['sku1'] == shipment
43+
assert order.allocation.sources['sku1'] == shipment
3744

3845

3946

@@ -44,7 +51,7 @@ def test_can_allocate_to_correct_shipment():
4451

4552
order.allocate(stock=Stock({}), shipments=[shipment1, shipment2])
4653

47-
assert order.allocation['sku2'] == shipment2
54+
assert order.allocation.sources['sku2'] == shipment2
4855

4956

5057
def test_allocates_to_stock_in_preference_to_shipment():
@@ -54,20 +61,20 @@ def test_allocates_to_stock_in_preference_to_shipment():
5461

5562
order.allocate(stock, shipments=[shipment])
5663

57-
assert order.allocation['sku1'] == stock
58-
assert stock['sku1'] == 990
59-
assert shipment['sku1'] == 1000
64+
assert order.allocation.sources['sku1'] == stock
65+
assert stock.quantities['sku1'] == 990
66+
assert shipment.quantities['sku1'] == 1000
6067

6168

6269
def test_can_allocate_multiple_lines_to_wh():
6370
order = Order({'sku1': 5, 'sku2': 10})
6471
stock = Stock({'sku1': 1000, 'sku2': 1000})
6572

6673
order.allocate(stock, shipments=[])
67-
assert order.allocation['sku1'] == stock
68-
assert order.allocation['sku2'] == stock
69-
assert stock['sku1'] == 995
70-
assert stock['sku2'] == 990
74+
assert order.allocation.sources['sku1'] == stock
75+
assert order.allocation.sources['sku2'] == stock
76+
assert stock.quantities['sku1'] == 995
77+
assert stock.quantities['sku2'] == 990
7178

7279

7380
def test_can_allocate_multiple_lines_to_shipment():
@@ -76,10 +83,10 @@ def test_can_allocate_multiple_lines_to_shipment():
7683

7784
order.allocate(stock=Stock({}), shipments=[shipment])
7885

79-
assert order.allocation['sku1'] == shipment
80-
assert order.allocation['sku2'] == shipment
81-
assert shipment['sku1'] == 995
82-
assert shipment['sku2'] == 990
86+
assert order.allocation.sources['sku1'] == shipment
87+
assert order.allocation.sources['sku2'] == shipment
88+
assert shipment.quantities['sku1'] == 995
89+
assert shipment.quantities['sku2'] == 990
8390

8491

8592
def test_can_allocate_to_both():
@@ -89,10 +96,10 @@ def test_can_allocate_to_both():
8996

9097
order.allocate(stock, shipments=[shipment])
9198

92-
assert order.allocation['sku1'] == stock
93-
assert order.allocation['sku2'] == shipment
94-
assert stock['sku1'] == 995
95-
assert shipment['sku2'] == 990
99+
assert order.allocation.sources['sku1'] == stock
100+
assert order.allocation.sources['sku2'] == shipment
101+
assert stock.quantities['sku1'] == 995
102+
assert shipment.quantities['sku2'] == 990
96103

97104

98105
def test_can_allocate_to_both_preferring_stock():
@@ -102,15 +109,15 @@ def test_can_allocate_to_both_preferring_stock():
102109

103110
order.allocate(stock, shipments=[shipment])
104111

105-
assert order.allocation['sku1'] == shipment
106-
assert order.allocation['sku2'] == shipment
107-
assert order.allocation['sku3'] == stock
108-
assert order.allocation['sku4'] == stock
109-
assert shipment['sku1'] == 999
110-
assert shipment['sku2'] == 998
111-
assert shipment['sku3'] == 1000
112-
assert stock['sku3'] == 997
113-
assert stock['sku4'] == 996
112+
assert order.allocation.sources['sku1'] == shipment
113+
assert order.allocation.sources['sku2'] == shipment
114+
assert order.allocation.sources['sku3'] == stock
115+
assert order.allocation.sources['sku4'] == stock
116+
assert shipment.quantities['sku1'] == 999
117+
assert shipment.quantities['sku2'] == 998
118+
assert shipment.quantities['sku3'] == 1000
119+
assert stock.quantities['sku3'] == 997
120+
assert stock.quantities['sku4'] == 996
114121

115122

116123
def test_mixed_allocation_are_avoided_if_possible():
@@ -120,8 +127,8 @@ def test_mixed_allocation_are_avoided_if_possible():
120127

121128
order.allocate(stock, shipments=[shipment])
122129

123-
assert order.allocation['sku1'] == shipment
124-
assert order.allocation['sku2'] == shipment
130+
assert order.allocation.sources['sku1'] == shipment
131+
assert order.allocation.sources['sku2'] == shipment
125132

126133

127134
def test_allocated_to_earliest_suitable_shipment_in_list():
@@ -132,8 +139,8 @@ def test_allocated_to_earliest_suitable_shipment_in_list():
132139

133140
order.allocate(stock, shipments=[shipment2, shipment1])
134141

135-
assert order.allocation['sku1'] == shipment1
136-
assert order.allocation['sku2'] == shipment1
142+
assert order.allocation.sources['sku1'] == shipment1
143+
assert order.allocation.sources['sku2'] == shipment1
137144

138145

139146
def test_still_chooses_earliest_if_split_across_shipments():
@@ -145,9 +152,9 @@ def test_still_chooses_earliest_if_split_across_shipments():
145152

146153
order.allocate(stock, shipments=[shipment3, shipment2, shipment1])
147154

148-
assert order.allocation['sku1'] == shipment1
149-
assert order.allocation['sku2'] == shipment2
150-
assert order.allocation['sku3'] == shipment2
155+
assert order.allocation.sources['sku1'] == shipment1
156+
assert order.allocation.sources['sku2'] == shipment2
157+
assert order.allocation.sources['sku3'] == shipment2
151158

152159

153160
def test_stock_not_quite_enough_means_we_use_shipment():
@@ -160,8 +167,8 @@ def test_stock_not_quite_enough_means_we_use_shipment():
160167

161168
order.allocate(stock, shipments=[shipment])
162169

163-
assert order.allocation['sku1'] == shipment
164-
assert order.allocation['sku2'] == shipment
170+
assert order.allocation.sources['sku1'] == shipment
171+
assert order.allocation.sources['sku2'] == shipment
165172

166173

167174
def test_cannot_allocate_if_insufficent_quantity_in_stock():
@@ -170,7 +177,7 @@ def test_cannot_allocate_if_insufficent_quantity_in_stock():
170177

171178
order.allocate(stock, shipments=[])
172179

173-
assert 'a-sku' not in order.allocation
180+
assert 'a-sku' not in order.allocation.skus
174181

175182

176183
def test_cannot_allocate_if_insufficent_quantity_in_shipment():
@@ -179,5 +186,5 @@ def test_cannot_allocate_if_insufficent_quantity_in_shipment():
179186

180187
order.allocate(stock=Stock({}), shipments=[shipment])
181188

182-
assert 'a-sku' not in order.allocation
189+
assert 'a-sku' not in order.allocation.skus
183190

0 commit comments

Comments
 (0)