diff --git a/jafgen/__init__.py b/jafgen/__init__.py index e69de29..e0227b5 100644 --- a/jafgen/__init__.py +++ b/jafgen/__init__.py @@ -0,0 +1,6 @@ +"""Jafgen, the Jaffle Shop data generator. + +Jafgen is a Python package that generates fake data that looks realistic +for Jaffle Shop, a fantasy coffee shop business. It does so by simulating +real customers attending the coffee shop in their daily lives. +""" diff --git a/jafgen/cli.py b/jafgen/cli.py index 44b9a25..46ab565 100644 --- a/jafgen/cli.py +++ b/jafgen/cli.py @@ -17,6 +17,7 @@ def run( typer.Option(help="Optional prefix for the output file names."), ] = "raw", ) -> None: + """Run jafgen in CLI mode.""" sim = Simulation(years, pre) sim.run_simulation() sim.save_results() diff --git a/jafgen/curves.py b/jafgen/curves.py index f4e24e0..34994f8 100644 --- a/jafgen/curves.py +++ b/jafgen/curves.py @@ -8,23 +8,33 @@ NumberArr = npt.NDArray[np.float64] | npt.NDArray[np.int32] - class Curve(ABC): + + """Base class for numerical curves that produce trends.""" + @property @abstractmethod def Domain(self) -> NumberArr: + """The function's domain, a.k.a the x-axis.""" raise NotImplementedError @abstractmethod def TranslateDomain(self, date: datetime.date) -> int: + """Translate the domain for a given date. + + This is useful for repeating series, like for making every + "sunday" on a weekly trend have the same value. + """ raise NotImplementedError @abstractmethod def Expr(self, x: float) -> float: + """Get `y` in `y=f(x)` where `f` is the trend.""" raise NotImplementedError @classmethod def eval(cls, date: datetime.date) -> float: + """Evaluate the curve's value (y-axis) on a given day.""" instance = cls() domain_value = instance.TranslateDomain(date) domain_index = domain_value % len(instance.Domain) @@ -33,6 +43,9 @@ def eval(cls, date: datetime.date) -> float: class AnnualCurve(Curve): + + """Produces trends over a year.""" + @property @override def Domain(self) -> NumberArr: @@ -48,13 +61,19 @@ def Expr(self, x: float) -> float: class WeekendCurve(Curve): + + """Produces trends over a weekend.""" + @property + @override def Domain(self) -> NumberArr: return np.array(range(6), dtype=np.float64) + @override def TranslateDomain(self, date: datetime.date) -> int: return date.weekday() - 1 + @override def Expr(self, x: float): if x >= 6: return 0.6 @@ -63,14 +82,19 @@ def Expr(self, x: float): class GrowthCurve(Curve): + + """Produces a growth over time trend.""" + @property + @override def Domain(self) -> NumberArr: return np.arange(500, dtype=np.int32) + @override def TranslateDomain(self, date: datetime.date) -> int: return (date.year - 2016) * 12 + date.month + @override def Expr(self, x: float) -> float: # ~ aim for ~20% growth/year return 1 + (x / 12) * 0.2 - diff --git a/jafgen/customers/customers.py b/jafgen/customers/customers.py index 9328756..d9c0364 100644 --- a/jafgen/customers/customers.py +++ b/jafgen/customers/customers.py @@ -5,6 +5,7 @@ import numpy as np from faker import Faker +from typing_extensions import override from jafgen.customers.order import Order from jafgen.customers.tweet import Tweet @@ -17,68 +18,71 @@ CustomerId = NewType("CustomerId", uuid.UUID) + @dataclass(frozen=True) class Customer(ABC): + + """Abstract base class for customers. + + A concrete implementation of `Customer` models the behavior of a customer. This + includes their probability to buy drinks/food, their probability to tweet about + their purchase etc. + """ + store: Store id: CustomerId = field(default_factory=lambda: CustomerId(fake.uuid4())) name: str = field(default_factory=fake.name) favorite_number: int = field(default_factory=lambda: fake.random.randint(1, 100)) fan_level: int = field(default_factory=lambda: fake.random.randint(1, 5)) - def p_buy_season(self, day: Day): - return self.store.p_buy(day) - def p_buy(self, day: Day) -> float: - p_buy_season = self.p_buy_season(day) + """Get the probability of buying something on a given day.""" + p_store_sell = self.store.p_sell(day) p_buy_persona = self.p_buy_persona(day) - p_buy_on_day = (p_buy_season * p_buy_persona) ** 0.5 + p_buy_on_day = (p_store_sell * p_buy_persona) ** 0.5 return p_buy_on_day @abstractmethod def p_buy_persona(self, day: Day) -> float: + """Get this customer's innate desire to buy something on a given day.""" raise NotImplementedError() @abstractmethod def p_tweet_persona(self, day: Day) -> float: + """Get this customer's innate desire to tweet about an order on a given day.""" raise NotImplementedError() def p_tweet(self, day: Day) -> float: + """Get the probability of tweeting about an order on a given day.""" return self.p_tweet_persona(day) - def get_order(self, day: Day) -> Order | None: + def get_order(self, day: Day) -> Order: + """Get this customer's order on a given day.""" items = self.get_order_items(day) order_minute = self.get_order_minute(day) order_day = day.at_minute(order_minute) - if not self.store.is_open_at(order_day): - return None - - return Order( - customer=self, - items=items, - store=self.store, - day=order_day - ) + return Order(customer=self, items=items, store=self.store, day=order_day) def get_tweet(self, order: Order) -> Tweet: + """Get this customer's tweet about an order.""" minutes_delta = int(fake.random.random() * 20) tweet_day = order.day.at_minute(order.day.total_minutes + minutes_delta) - return Tweet( - customer=self, - order=order, - day=tweet_day - ) + return Tweet(customer=self, order=order, day=tweet_day) @abstractmethod def get_order_items(self, day: Day) -> list[Item]: + """Get the list of ordered items on a given day.""" raise NotImplementedError() @abstractmethod def get_order_minute(self, day: Day) -> int: + """Get the time the customer decided to order on a given day.""" raise NotImplementedError() - def sim_day(self, day: Day): + def sim_day(self, day: Day) -> tuple[Order | None, Tweet | None]: + """Simulate a day in the life of this customer.""" p_buy = self.p_buy(day) p_buy_threshold = np.random.random() p_tweet = self.p_tweet(day) @@ -86,7 +90,7 @@ def sim_day(self, day: Day): if p_buy > p_buy_threshold: if p_tweet > p_tweet_threshold: order = self.get_order(day) - if order and len(order.items) > 0: + if self.store.is_open(order.day) and len(order.items) > 0: return order, self.get_tweet(order) else: return None, None @@ -96,6 +100,10 @@ def sim_day(self, day: Day): return None, None def to_dict(self) -> dict[str, Any]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "id": str(self.id), "name": str(self.name), @@ -103,15 +111,19 @@ def to_dict(self) -> dict[str, Any]: class RemoteWorker(Customer): - """This person works from a coffee shop""" + """Pretending to work while staring at their Macbook full of stickers.""" + + @override def p_buy_persona(self, day: Day): buy_propensity = (self.favorite_number / 100) * 0.4 return 0.001 if day.is_weekend else buy_propensity + @override def p_tweet_persona(self, day: Day): return 0.01 + @override def get_order_minute(self, day: Day) -> int: # most likely to order in the morning # exponentially less likely to order in the afternoon @@ -119,6 +131,7 @@ def get_order_minute(self, day: Day) -> int: order_time = np.random.normal(loc=avg_time, scale=180) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): num_drinks = 1 food = [] @@ -133,36 +146,47 @@ def get_order_items(self, day: Day): class BrunchCrowd(Customer): - """Do you sell mimosas?""" + """Do you sell mimosas?.""" + + @override def p_buy_persona(self, day: Day): buy_propensity = 0.2 + (self.favorite_number / 100) * 0.2 return buy_propensity if day.is_weekend else 0 + @override def p_tweet_persona(self, day: Day): return 0.8 + @override def get_order_minute(self, day: Day) -> int: # most likely to order in the early afternoon avg_time = 300 + ((self.favorite_number - 50) / 50) * 120 order_time = np.random.normal(loc=avg_time, scale=120) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): num_customers = 1 + int(self.favorite_number / 20) - return Inventory.get_item_type(ItemType.JAFFLE, num_customers) + Inventory.get_item_type(ItemType.BEVERAGE, num_customers) + return Inventory.get_item_type( + ItemType.JAFFLE, num_customers + ) + Inventory.get_item_type(ItemType.BEVERAGE, num_customers) class Commuter(Customer): - """the regular, thanks""" + """The regular, thanks.""" + + @override def p_buy_persona(self, day: Day): buy_propensity = 0.5 + (self.favorite_number / 100) * 0.3 return 0.001 if day.is_weekend else buy_propensity + @override def p_tweet_persona(self, day: Day): return 0.2 + @override def get_order_minute(self, day: Day) -> int: # most likely to order in the morning # exponentially less likely to order in the afternoon @@ -170,13 +194,16 @@ def get_order_minute(self, day: Day) -> int: order_time = np.random.normal(loc=avg_time, scale=30) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): return Inventory.get_item_type(ItemType.BEVERAGE, 1) class Student(Customer): - """coffee might help""" + """Coffee might help.""" + + @override def p_buy_persona(self, day: Day): if day.season == Season.SUMMER: return 0 @@ -184,15 +211,18 @@ def p_buy_persona(self, day: Day): buy_propensity = 0.1 + (self.favorite_number / 100) * 0.4 return buy_propensity + @override def p_tweet_persona(self, day: Day): return 0.8 + @override def get_order_minute(self, day: Day) -> int: # later is better avg_time = 9 * 60 order_time = np.random.normal(loc=avg_time, scale=120) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): food = [] if fake.random.random() > 0.5: @@ -202,41 +232,53 @@ def get_order_items(self, day: Day): class Casuals(Customer): - """just popping in""" + """Just popping in.""" + + @override def p_buy_persona(self, day: Day): return 0.1 + @override def p_tweet_persona(self, day: Day): return 0.1 + @override def get_order_minute(self, day: Day) -> int: avg_time = 5 * 60 order_time = np.random.normal(loc=avg_time, scale=120) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): num_drinks = int(fake.random.random() * 10 / 3) num_food = int(fake.random.random() * 10 / 3) - return Inventory.get_item_type(ItemType.BEVERAGE, num_drinks) + Inventory.get_item_type(ItemType.JAFFLE, num_food) + return Inventory.get_item_type( + ItemType.BEVERAGE, num_drinks + ) + Inventory.get_item_type(ItemType.JAFFLE, num_food) class HealthNut(Customer): - """A light beverage in the sunshine as a treat""" + """A light beverage in the sunshine as a treat.""" + + @override def p_buy_persona(self, day: Day): if day.season == Season.SUMMER: buy_propensity = 0.1 + (self.favorite_number / 100) * 0.4 return buy_propensity return 0.2 + @override def p_tweet_persona(self, day: Day): return 0.6 + @override def get_order_minute(self, day: Day) -> int: avg_time = 5 * 60 order_time = np.random.normal(loc=avg_time, scale=120) return max(0, int(order_time)) + @override def get_order_items(self, day: Day): return Inventory.get_item_type(ItemType.BEVERAGE, 1) diff --git a/jafgen/customers/order.py b/jafgen/customers/order.py index 5d59e87..43de36f 100644 --- a/jafgen/customers/order.py +++ b/jafgen/customers/order.py @@ -13,8 +13,12 @@ OrderId = NewType("OrderId", uuid.UUID) + @dataclass class Order: + + """An order of a few items from a single customer at a store.""" + customer: "customer.Customer" day: Day store: Store @@ -26,14 +30,20 @@ class Order: total: float = field(init=False) def __post_init__(self) -> None: + """Initialize subtotal, tax_paid and total based on the items.""" self.subtotal = sum(i.price for i in self.items) self.tax_paid = self.store.tax_rate * self.subtotal self.total = self.subtotal + self.tax_paid def __str__(self): + """Get a human readable string that represents this order.""" return f"{self.customer.name} bought {str(self.items)} at {self.day}" def to_dict(self) -> dict[str, Any]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "id": str(self.id), "customer": str(self.customer.id), @@ -47,5 +57,3 @@ def to_dict(self) -> dict[str, Any]: "order_total": int(int(self.subtotal * 100) + int(self.tax_paid * 100)), } - def items_to_dict(self) -> list[dict[str, Any]]: - return [item.to_dict() for item in self.items] diff --git a/jafgen/customers/tweet.py b/jafgen/customers/tweet.py index d4565c0..6798338 100644 --- a/jafgen/customers/tweet.py +++ b/jafgen/customers/tweet.py @@ -15,6 +15,9 @@ @dataclass class Tweet: + + """A tweet created by a customer after ordering something.""" + day: Day customer: "customer.Customer" order: Order @@ -22,9 +25,14 @@ class Tweet: content: str = field(init=False) def __post_init__(self) -> None: + """Lazily initialize the contents of this tweet.""" self.content = self._construct_tweet() def to_dict(self) -> dict[str, str]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "id": str(self.id), "user_id": str(self.customer.id), @@ -36,9 +44,16 @@ def _construct_tweet(self) -> str: if len(self.order.items) == 1: items_sentence = f"Ordered a {self.order.items[0].name}" elif len(self.order.items) == 2: - items_sentence = f"Ordered a {self.order.items[0].name} and a {self.order.items[1].name}" + items_sentence = ( + f"Ordered a {self.order.items[0].name} and a {self.order.items[1].name}" + ) else: - items_sentence = f"Ordered a {', a '.join(item.name for item in self.order.items[:-1])}, and a {self.order.items[-1].name}" + items_sentence = ( + "Ordered a " + + ", a ".join(item.name for item in self.order.items[:-1]) + + ", and a " + + self.order.items[-1].name + ) if self.customer.fan_level > 3: adjective = fake.random.choice( [ diff --git a/jafgen/simulation.py b/jafgen/simulation.py index 873f9b4..0f2a12b 100644 --- a/jafgen/simulation.py +++ b/jafgen/simulation.py @@ -23,8 +23,13 @@ T_3PM = time_from_total_minutes(60 * 15) T_8PM = time_from_total_minutes(60 * 20) + class Simulation: + + """Runs a simulation of multiple days of our customers' lives.""" + def __init__(self, years: int, prefix: str): + """Initialize the simulation.""" self.years = years self.scale = 100 self.prefix = prefix @@ -61,6 +66,7 @@ def __init__(self, years: int, prefix: str): self.sim_days = 365 * self.years def run_simulation(self): + """Run the simulation.""" for i in track( range(self.sim_days), description="🥪 Pressing fresh jaffles..." ): @@ -75,6 +81,7 @@ def run_simulation(self): self.tweets.append(tweet) def save_results(self) -> None: + """Save the simulated results to `./jaffle-data/[prefix]_[entity].csv`.""" stock: Stock = Stock() inventory: Inventory = Inventory() entities: dict[str, list[dict[str, Any]]] = { diff --git a/jafgen/stores/inventory.py b/jafgen/stores/inventory.py index 009054f..13b2536 100644 --- a/jafgen/stores/inventory.py +++ b/jafgen/stores/inventory.py @@ -9,21 +9,30 @@ class Inventory: + + """Holds all possible items to buy from a Jaffle Shop store.""" + inventory: dict[ItemType, list[Item]] = {} @classmethod def update(cls, inventory_list: list[Item]): - cls.inventory[ItemType.JAFFLE]= [] + """Update the inventory with a new list of items.""" + cls.inventory[ItemType.JAFFLE] = [] cls.inventory[ItemType.BEVERAGE] = [] for item in inventory_list: cls.inventory[item.type].append(item) @classmethod def get_item_type(cls, type: ItemType, count: int = 1): - return [fake.random.choice(cls.inventory[type])for _ in range(count)] + """Get a random assortment of N items of `type`.""" + return [fake.random.choice(cls.inventory[type]) for _ in range(count)] @classmethod def to_dict(cls) -> list[dict[str, Any]]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ all_items: list[dict[str, Any]] = [] for key in cls.inventory: all_items += [item.to_dict() for item in cls.inventory[key]] @@ -49,19 +58,23 @@ def to_dict(cls) -> list[dict[str, Any]]: Item( sku=SKU("JAF-003"), name="the krautback", - description="lamb and pork bratwurst with house-pickled cabbage sauerkraut and mustard", + description="lamb and pork bratwurst with house-pickled " + "cabbage sauerkraut and mustard", type=ItemType.JAFFLE, price=12, ), Item( sku=SKU("JAF-004"), name="flame impala", - description="pulled pork and pineapple al pastor marinated in ghost pepper sauce, kevin parker's favorite! ", + description="pulled pork and pineapple al pastor marinated " + "in ghost pepper sauce, kevin parker's favorite! ", type=ItemType.JAFFLE, price=14, ), Item( - sku=SKU("JAF-005"), name="mel-bun", description="melon and minced beef bao, in a jaffle, savory and sweet", + sku=SKU("JAF-005"), + name="mel-bun", + description="melon and minced beef bao, in a jaffle, savory and sweet", type=ItemType.JAFFLE, price=12, ), @@ -89,14 +102,16 @@ def to_dict(cls) -> list[dict[str, Any]]: Item( sku=SKU("BEV-004"), name="for richer or pourover ", - description="daily selection of single estate beans for a delicious hot pourover", + description="daily selection of single estate beans " + "for a delicious hot pourover", type=ItemType.BEVERAGE, price=7, ), Item( sku=SKU("BEV-005"), name="adele-ade", - description="a kiwi and lime agua fresca, hello from the other side of thirst", + description="a kiwi and lime agua fresca, hello from the " + "other side of thirst", type=ItemType.BEVERAGE, price=4, ), diff --git a/jafgen/stores/item.py b/jafgen/stores/item.py index 11e0b6a..beeaf2e 100644 --- a/jafgen/stores/item.py +++ b/jafgen/stores/item.py @@ -6,25 +6,32 @@ class ItemType(str, Enum): + + """The type of an item (food or drink).""" + JAFFLE = "JAFFLE" BEVERAGE = "BEVERAGE" @dataclass(frozen=True) class Item: + + """Anything that can be sold on a Jaffle Shop store. + + Can be coffee, smoothies, sandwiches etc. + """ + sku: StorageKeepingUnit name: str description: str type: ItemType price: float - def __str__(self): - return f"<{self.name} @ ${self.price}>" - - def __repr__(self): - return self.__str__() - def to_dict(self) -> dict[str, Any]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "sku": self.sku, "name": str(self.name), diff --git a/jafgen/stores/market.py b/jafgen/stores/market.py index 44629bc..b1991bd 100644 --- a/jafgen/stores/market.py +++ b/jafgen/stores/market.py @@ -19,7 +19,11 @@ fake = Faker() + class Market: + + """A bunch of people together buying from a store.""" + PersonaMix = [ (Commuter, 0.25), (RemoteWorker, 0.25), @@ -29,7 +33,19 @@ class Market: (HealthNut, 0.1), ] - def __init__(self, store: Store, num_customers: int, days_to_penetration: int = 365): + def __init__( + self, store: Store, num_customers: int, days_to_penetration: int = 365 + ): + """Initialize the market. + + Args: + ---- + store: the store that fulfills this market's desires. + num_customers: number of customers in this market. + days_to_penetration: number of days until this store reaches full market + penetration. + + """ self.store = store self.num_customers = num_customers self.days_to_penetration = days_to_penetration @@ -45,6 +61,7 @@ def __init__(self, store: Store, num_customers: int, days_to_penetration: int = fake.random.shuffle(self.addressable_customers) def sim_day(self, day: Day) -> Iterator[tuple[Order | None, Tweet | None]]: + """Simulate a day in this market.""" days_since_open = self.store.days_since_open(day) if days_since_open < 0: yield None, None diff --git a/jafgen/stores/stock.py b/jafgen/stores/stock.py index 32193ae..dd0420e 100644 --- a/jafgen/stores/stock.py +++ b/jafgen/stores/stock.py @@ -5,10 +5,14 @@ class Stock: + + """Holds supplies for Jaffle Shop stores.""" + stock: dict[SKU, list[Supply]] = {} @classmethod def update(cls, stock_list: list[Supply]): + """Update the items available in stock.""" for supply in stock_list: skus = supply.skus for sku in skus: @@ -18,11 +22,16 @@ def update(cls, stock_list: list[Supply]): @classmethod def to_dict(cls) -> list[dict[str, Any]]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ all_items: list[dict[str, Any]] = [] for key in cls.stock: all_items += [item.to_dict(key) for item in cls.stock[key]] return all_items + Stock.update( [ Supply( @@ -30,59 +39,111 @@ def to_dict(cls) -> list[dict[str, Any]]: name="compostable cutlery - knife", cost=0.07, perishable=False, - skus=[SKU("JAF-001"), SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], + skus=[ + SKU("JAF-001"), + SKU("JAF-002"), + SKU("JAF-003"), + SKU("JAF-004"), + SKU("JAF-005"), + ], ), Supply( id=SupplyId("SUP-002"), name="cutlery - fork", cost=0.07, perishable=False, - skus=[SKU("JAF-001"), SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], + skus=[ + SKU("JAF-001"), + SKU("JAF-002"), + SKU("JAF-003"), + SKU("JAF-004"), + SKU("JAF-005"), + ], ), Supply( id=SupplyId("SUP-003"), name="serving boat", cost=0.11, perishable=False, - skus=[SKU("JAF-001"), SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], + skus=[ + SKU("JAF-001"), + SKU("JAF-002"), + SKU("JAF-003"), + SKU("JAF-004"), + SKU("JAF-005"), + ], ), Supply( id=SupplyId("SUP-004"), name="napkin", cost=0.04, perishable=False, - skus=[SKU("JAF-001"), SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], + skus=[ + SKU("JAF-001"), + SKU("JAF-002"), + SKU("JAF-003"), + SKU("JAF-004"), + SKU("JAF-005"), + ], ), Supply( id=SupplyId("SUP-005"), name="16oz compostable clear cup", cost=0.13, perishable=False, - skus=[SKU("BEV-001"), SKU("BEV-002"), SKU("BEV-003"), SKU("BEV-004"), SKU("BEV-005")], + skus=[ + SKU("BEV-001"), + SKU("BEV-002"), + SKU("BEV-003"), + SKU("BEV-004"), + SKU("BEV-005"), + ], ), Supply( id=SupplyId("SUP-006"), name="16oz compostable clear lid", cost=0.04, perishable=False, - skus=[SKU("BEV-001"), SKU("BEV-002"), SKU("BEV-003"), SKU("BEV-004"), SKU("BEV-005")], + skus=[ + SKU("BEV-001"), + SKU("BEV-002"), + SKU("BEV-003"), + SKU("BEV-004"), + SKU("BEV-005"), + ], ), Supply( id=SupplyId("SUP-007"), name="biodegradable straw", cost=0.13, perishable=False, - skus=[SKU("BEV-001"), SKU("BEV-002"), SKU("BEV-003"), SKU("BEV-004"), SKU("BEV-005")], + skus=[ + SKU("BEV-001"), + SKU("BEV-002"), + SKU("BEV-003"), + SKU("BEV-004"), + SKU("BEV-005"), + ], ), Supply( - id=SupplyId("SUP-008"), name="chai mix", cost=0.98, perishable=True, skus=[SKU("BEV-002")] + id=SupplyId("SUP-008"), + name="chai mix", + cost=0.98, + perishable=True, + skus=[SKU("BEV-002")], ), Supply( id=SupplyId("SUP-009"), name="bread", cost=0.33, perishable=True, - skus=[SKU("JAF-001"), SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], + skus=[ + SKU("JAF-001"), + SKU("JAF-002"), + SKU("JAF-003"), + SKU("JAF-004"), + SKU("JAF-005"), + ], ), Supply( id=SupplyId("SUP-010"), @@ -92,13 +153,25 @@ def to_dict(cls) -> list[dict[str, Any]]: skus=[SKU("JAF-002"), SKU("JAF-003"), SKU("JAF-004"), SKU("JAF-005")], ), Supply( - id=SupplyId("SUP-011"), name="nutella", cost=0.46, perishable=True, skus=[SKU("JAF-001")] + id=SupplyId("SUP-011"), + name="nutella", + cost=0.46, + perishable=True, + skus=[SKU("JAF-001")], ), Supply( - id=SupplyId("SUP-012"), name="banana", cost=0.13, perishable=True, skus=[SKU("JAF-001")] + id=SupplyId("SUP-012"), + name="banana", + cost=0.13, + perishable=True, + skus=[SKU("JAF-001")], ), Supply( - id=SupplyId("SUP-013"), name="beef stew", cost=1.69, perishable=True, skus=[SKU("JAF-002")] + id=SupplyId("SUP-013"), + name="beef stew", + cost=1.69, + perishable=True, + skus=[SKU("JAF-002")], ), Supply( id=SupplyId("SUP-014"), @@ -115,7 +188,11 @@ def to_dict(cls) -> list[dict[str, Any]]: skus=[SKU("JAF-003")], ), Supply( - id=SupplyId("SUP-016"), name="mustard", cost=0.07, perishable=True, skus=[SKU("JAF-003")] + id=SupplyId("SUP-016"), + name="mustard", + cost=0.07, + perishable=True, + skus=[SKU("JAF-003")], ), Supply( id=SupplyId("SUP-017"), @@ -125,10 +202,18 @@ def to_dict(cls) -> list[dict[str, Any]]: skus=[SKU("JAF-004")], ), Supply( - id=SupplyId("SUP-018"), name="pineapple", cost=0.26, perishable=True, skus=[SKU("JAF-004")] + id=SupplyId("SUP-018"), + name="pineapple", + cost=0.26, + perishable=True, + skus=[SKU("JAF-004")], ), Supply( - id=SupplyId("SUP-019"), name="melon", cost=0.33, perishable=True, skus=[SKU("JAF-005")] + id=SupplyId("SUP-019"), + name="melon", + cost=0.33, + perishable=True, + skus=[SKU("JAF-005")], ), Supply( id=SupplyId("SUP-020"), @@ -145,13 +230,25 @@ def to_dict(cls) -> list[dict[str, Any]]: skus=[SKU("JAF-004")], ), Supply( - id=SupplyId("SUP-022"), name="mango", cost=0.32, perishable=True, skus=[SKU("BEV-001")] + id=SupplyId("SUP-022"), + name="mango", + cost=0.32, + perishable=True, + skus=[SKU("BEV-001")], ), Supply( - id=SupplyId("SUP-023"), name="tangerine", cost=0.2, perishable=True, skus=[SKU("BEV-001")] + id=SupplyId("SUP-023"), + name="tangerine", + cost=0.2, + perishable=True, + skus=[SKU("BEV-001")], ), Supply( - id=SupplyId("SUP-024"), name="oatmilk", cost=0.11, perishable=True, skus=[SKU("BEV-002")] + id=SupplyId("SUP-024"), + name="oatmilk", + cost=0.11, + perishable=True, + skus=[SKU("BEV-002")], ), Supply( id=SupplyId("SUP-025"), @@ -174,7 +271,19 @@ def to_dict(cls) -> list[dict[str, Any]]: perishable=True, skus=[SKU("BEV-003")], ), - Supply(id=SupplyId("SUP-028"), name="kiwi", cost=0.2, perishable=True, skus=[SKU("BEV-005")]), - Supply(id=SupplyId("SUP-029"), name="lime", cost=0.13, perishable=True, skus=[SKU("BEV-005")]), + Supply( + id=SupplyId("SUP-028"), + name="kiwi", + cost=0.2, + perishable=True, + skus=[SKU("BEV-005")], + ), + Supply( + id=SupplyId("SUP-029"), + name="lime", + cost=0.13, + perishable=True, + skus=[SKU("BEV-005")], + ), ] ) diff --git a/jafgen/stores/store.py b/jafgen/stores/store.py index 1b52dc3..84cc4aa 100644 --- a/jafgen/stores/store.py +++ b/jafgen/stores/store.py @@ -1,7 +1,6 @@ -import datetime as dt import uuid from dataclasses import dataclass, field -from typing import Iterator, NewType +from typing import NewType from faker import Faker @@ -11,8 +10,12 @@ StoreId = NewType("StoreId", uuid.UUID) + @dataclass(frozen=True) class Store: + + """A single Jaffle Shop store.""" + name: str base_popularity: float hours_of_operation: WeekHoursOfOperation @@ -20,31 +23,25 @@ class Store: tax_rate: float id: StoreId = field(default_factory=lambda: StoreId(fake.uuid4())) - def p_buy(self, day: Day) -> float: + def p_sell(self, day: Day) -> float: + """Get the probability of this store selling something on a given day.""" return self.base_popularity * day.get_effect() - def minutes_open(self, day: Day) -> int: - return self.hours_of_operation.total_minutes_open(day) - - def iter_minutes_open(self, day: Day) -> Iterator[int]: - yield from self.hours_of_operation.iter_minutes(day) - def is_open(self, day: Day) -> bool: - return day.date >= self.opened_day.date - - def is_open_at(self, day: Day) -> bool: + """Whether the store is open on a given day.""" + if day.date < self.opened_day.date: + return False return self.hours_of_operation.is_open(day) def days_since_open(self, day: Day) -> int: + """Get the number of days since this store has opened.""" return day.date_index - self.opened_day.date_index - def opens_at(self, day: Day) -> dt.time: - return self.hours_of_operation.opens_at(day) - - def closes_at(self, day: Day) -> dt.time: - return self.hours_of_operation.closes_at(day) - def to_dict(self) -> dict[str, str]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "id": str(self.id), "name": str(self.name), diff --git a/jafgen/stores/supply.py b/jafgen/stores/supply.py index bfb60b6..c4a8fd8 100644 --- a/jafgen/stores/supply.py +++ b/jafgen/stores/supply.py @@ -4,21 +4,26 @@ SupplyId = NewType("SupplyId", str) StorageKeepingUnit = NewType("StorageKeepingUnit", str) + @dataclass(frozen=True) class Supply: + + """Any supply that's not an item itself. + + Examples: cuttlery, ingredients, napkins etc. + """ + id: SupplyId name: str cost: float perishable: bool skus: list[StorageKeepingUnit] - def __str__(self): - return f"<{self.name} @ ${self.cost}>" - - def __repr__(self): - return self.__str__() - def to_dict(self, sku: StorageKeepingUnit) -> dict[str, str | int]: + """Serialize to dict. + + TODO: replace this by serializer class. + """ return { "id": str(self.id), "name": str(self.name), diff --git a/jafgen/time.py b/jafgen/time.py index bf3692c..41283ec 100644 --- a/jafgen/time.py +++ b/jafgen/time.py @@ -1,5 +1,5 @@ import datetime as dt -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Iterator @@ -36,6 +36,9 @@ def total_minutes_elapsed(t: dt.time | dt.timedelta) -> int: class Season(str, Enum): + + """A season of the year.""" + WINTER = "WINTER" SPRING = "SPRING" SUMMER = "SUMMER" @@ -43,6 +46,7 @@ class Season(str, Enum): @classmethod def from_date(cls, date: dt.date) -> "Season": + """Get the date's Season (in the northern hemisphere).""" month_no = date.month day_no = date.day @@ -57,26 +61,41 @@ def from_date(cls, date: dt.date) -> "Season": return cls.WINTER -@dataclass(init=False) + +@dataclass(frozen=True) class Day: + + """A wrapper around `datetime.date` with some convenience functions.""" + EPOCH = dt.datetime(year=2018, month=9, day=1) SEASONAL_MONTHLY_CURVE = AnnualCurve() WEEKEND_CURVE = WeekendCurve() GROWTH_CURVE = GrowthCurve() - def __init__(self, date_index: int, minutes: int = 0): - self.date_index = date_index - self.date = self.EPOCH + dt.timedelta(days=date_index, minutes=minutes) - self.effects = [ + date_index: int + minutes: int = 0 + date: dt.datetime = field(init=False) + effects: list[float] = field(init=False) + + def __post_init__(self): + """Set `date` and `effects` based on `__init__` args.""" + # Have to use object.__setattr__ since it's frozen dataclass + object.__setattr__(self, "date", self.EPOCH + dt.timedelta( + days=self.date_index, + minutes=self.minutes) + ) + object.__setattr__(self, "effects", [ self.SEASONAL_MONTHLY_CURVE.eval(self.date), self.WEEKEND_CURVE.eval(self.date), self.GROWTH_CURVE.eval(self.date), - ] + ]) def at_minute(self, minutes: int) -> "Day": + """Get a new instance of `Day` in the same day but different minute.""" return Day(self.date_index, minutes=minutes) def get_effect(self) -> float: + """Get a simulation effect multiplier based on all of this day's effects.""" total = 1 for effect in self.effects: total = total * effect @@ -84,40 +103,58 @@ def get_effect(self) -> float: @property def day_of_week(self) -> int: + """Get the day of the week, where Monday=0 and Sunday=6.""" return self.date.weekday() @property def is_weekend(self) -> bool: - # 5 + 6 are weekends + """Whether this day is in the weekend.""" return self.date.weekday() >= 5 @property def season(self) -> Season: + """Get the season of the year this day is in.""" return Season.from_date(self.date) @property def total_minutes(self) -> int: + """Get the total elapsed minutes of this day.""" return self.date.hour * 60 + self.date.minute + @dataclass(frozen=True) class DayHoursOfOperation: + + """A shop's hours of operations during a single day.""" + opens_at: dt.time closes_at: dt.time @property def total_minutes_open(self) -> int: + """The total minutes the shop will be open this day.""" time_open = time_delta_sub(self.closes_at, time_to_delta(self.opens_at)) return total_minutes_elapsed(time_open) def is_open(self, time: dt.time) -> bool: + """Whether the shop will be open during `time`.""" return time >= self.opens_at and time < self.closes_at def iter_minutes(self) -> Iterator[int]: + """Iterate over all the minutes in this day of operation.""" for minute in range(self.total_minutes_open): yield minute + @dataclass(frozen=True) class WeekHoursOfOperation: + + """A shop's hours of operations during a week. + + Consists of two collections of `DayHoursOfOperation`: one for the week days + and another one for the weekends. + """ + week_days: DayHoursOfOperation weekends: DayHoursOfOperation @@ -125,18 +162,22 @@ def _get_todays_schedule(self, day: Day) -> DayHoursOfOperation: return self.weekends if day.is_weekend else self.week_days def opens_at(self, day: Day) -> dt.time: + """Get the time the shop will open on `day`.""" return self._get_todays_schedule(day).opens_at def closes_at(self, day: Day) -> dt.time: + """Get the time the shop will close on `day`.""" return self._get_todays_schedule(day).closes_at def total_minutes_open(self, day: Day) -> int: + """Get the total minutes the shop will be open on `day`.""" return self._get_todays_schedule(day).total_minutes_open def is_open(self, day: Day) -> bool: + """Whether the shop is open on `day`.""" time = time_from_total_minutes(day.total_minutes) return self._get_todays_schedule(day).is_open(time) def iter_minutes(self, day: Day) -> Iterator[int]: + """Iterate over all the minutes in this day of operation.""" yield from self._get_todays_schedule(day).iter_minutes() - diff --git a/lefthook.yaml b/lefthook.yaml index f1ba17d..618d9b0 100644 --- a/lefthook.yaml +++ b/lefthook.yaml @@ -1,7 +1,9 @@ pre-commit: commands: - # ruff: - # run: .venv/bin/ruff check --fix + ruff-format: + run: .venv/bin/ruff format + ruff-lint: + run: .venv/bin/ruff check --fix pyright: run: .venv/bin/pyright diff --git a/pyproject.toml b/pyproject.toml index 8dd44be..92bbf97 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,12 @@ select = [ "T20", # Print statements "I", # isort ] +ignore = [ + "D100", # Doc in public modules + "D104", + "D211", # This is incompatible with D203 + "D213", # This is incompatible with D212 +] # Allow autofix for all enabled rules (when `--fix`) is provided. fixable = [ "A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", @@ -35,5 +41,9 @@ fixable = [ "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT" ] +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["D"] +"setup.py" = ["D"] + [tool.pydocstyle] convention = "google" diff --git a/tests/conftest.py b/tests/conftest.py index 2c0e595..e016920 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,7 @@ T_3PM = time_from_total_minutes(60 * 15) T_8PM = time_from_total_minutes(60 * 20) + @pytest.fixture def default_store() -> Store: """Return a pre-initialized store that can be used for tests.""" diff --git a/tests/test_order_totals.py b/tests/test_order_totals.py index 90c0d32..dd52126 100644 --- a/tests/test_order_totals.py +++ b/tests/test_order_totals.py @@ -1,9 +1,9 @@ -from jafgen.time import Day -from jafgen.customers.customers import Customer, BrunchCrowd, RemoteWorker, Student +from jafgen.customers.customers import BrunchCrowd, Customer, RemoteWorker, Student from jafgen.customers.order import Order -from jafgen.stores.item import ItemType from jafgen.stores.inventory import Inventory +from jafgen.stores.item import ItemType from jafgen.stores.store import Store +from jafgen.time import Day def test_order_totals(default_store: Store): @@ -16,9 +16,8 @@ def test_order_totals(default_store: Store): orders.append( Order( customer=CustType(store=default_store), - items= - inventory.get_item_type(ItemType.JAFFLE, 2) + - inventory.get_item_type(ItemType.BEVERAGE, 1), + items=inventory.get_item_type(ItemType.JAFFLE, 2) + + inventory.get_item_type(ItemType.BEVERAGE, 1), store=default_store, day=Day(date_index=i), ) @@ -27,15 +26,13 @@ def test_order_totals(default_store: Store): for order in orders: assert ( order.subtotal - == order.items[0].price - + order.items[1].price - + order.items[2].price + == order.items[0].price + order.items[1].price + order.items[2].price ) assert order.tax_paid == order.subtotal * order.store.tax_rate assert order.total == order.subtotal + order.tax_paid - assert round(float(order.total), 2) == round( - float(order.subtotal), 2 - ) + round(float(order.tax_paid), 2) + assert round(float(order.total), 2) == round(float(order.subtotal), 2) + round( + float(order.tax_paid), 2 + ) order_dict = order.to_dict() assert ( order_dict["order_total"] == order_dict["subtotal"] + order_dict["tax_paid"] diff --git a/tests/test_tweets.py b/tests/test_tweets.py index dfee52c..28acf1a 100644 --- a/tests/test_tweets.py +++ b/tests/test_tweets.py @@ -1,4 +1,3 @@ -from jafgen.time import Day from jafgen.customers.customers import ( BrunchCrowd, Casuals, @@ -9,6 +8,7 @@ Student, ) from jafgen.stores.store import Store +from jafgen.time import Day def test_tweets(default_store: Store): @@ -27,7 +27,9 @@ def test_tweets(default_store: Store): assert tweet.order == order assert ( tweet.day.date - <= tweet.order.day.at_minute(tweet.order.day.total_minutes + 20).date + <= tweet.order.day.at_minute( + tweet.order.day.total_minutes + 20 + ).date ) if not order: assert not tweet