diff --git a/jafgen/simulation.py b/jafgen/simulation.py index a6a87a3..4a5a94c 100644 --- a/jafgen/simulation.py +++ b/jafgen/simulation.py @@ -107,13 +107,13 @@ def run_simulation(self): ): for market in self.markets: day = Day(i) - for result in market.sim_day(day): - if result is not None and all(val is not None for val in result): - order, tweet = result + for order, tweet in market.sim_day(day): + if order: self.orders.append(order) - self.tweets.append(tweet) if order.customer.customer_id not in self.customers: self.customers[order.customer.customer_id] = order.customer + if tweet: + self.tweets.append(tweet) def save_results(self) -> None: stock: Stock = Stock() diff --git a/jafgen/stores/market.py b/jafgen/stores/market.py index 740c027..792c5cc 100644 --- a/jafgen/stores/market.py +++ b/jafgen/stores/market.py @@ -1,4 +1,5 @@ import random +from typing import Iterator import numpy as np @@ -6,10 +7,12 @@ BrunchCrowd, Casuals, Commuter, + HealthNut, RemoteWorker, Student, - HealthNut, ) +from jafgen.customers.order import Order +from jafgen.customers.tweet import Tweet class Market(object): @@ -38,10 +41,10 @@ def __init__(self, store, num_customers, days_to_penetration=365): self.active_customers = [] - def sim_day(self, day): + def sim_day(self, day) -> Iterator[tuple[Order | None, Tweet | None]]: days_since_open = self.store.days_since_open(day) if days_since_open < 0: - yield None + yield None, None return elif days_since_open < 7: pct_penetration = min(days_since_open / self.days_to_penetration, 1) @@ -53,7 +56,7 @@ def sim_day(self, day): num_desired_customers = market_penetration * len(self.addressable_customers) customers_to_add = int(num_desired_customers - len(self.active_customers)) - for i in range(customers_to_add): + for _ in range(customers_to_add): customer = self.addressable_customers.pop() self.active_customers.append(customer)