Added option to have days (#124)

README.md

@@ -39,6 +39,8 @@ pip install jafgen
 
 The following options are available:
 
+- `--days [int]` The number of days to generate data for. If both years and days are set, they will be added together.
+
 - `--pre` sets a prefix for the generated files in the format `[prefix]_[file_name].csv`. It defaults to `raw`.
 
 Generate a simulation spanning 3 years from 2016-2019 with a prefix of `cool`:

jafgen/cli.py

@@ -9,14 +9,24 @@
 
 @app.command()
 def run(
+    # We set default to 0 here to make sure users don't get confused if they only put in days.
+    # If they don't set anything we have years default = 1 later to keep the same functionality. 
     years: Annotated[
-        int, typer.Argument(help="Number of years to simulate. Default is 1.")
-    ] = 1,
+         int, typer.Argument(help="Number of years to simulate. If neither days nor years are provided, the default is 1 year.")
+    ] = 0,
+    days: Annotated[
+        int, typer.Option(help="Number of days to simulate. Default is 0.")
+    ] = 0,
     pre: Annotated[
         str,
         typer.Option(help="Optional prefix for the output file names."),
     ] = "raw",
 ) -> None:
-    sim = Simulation(years, pre)
+
+    # To keep the default value for backwards compatibility.
+    if years == 0 and days == 0:
+        years = 1
+
+    sim = Simulation(years, days, pre)
     sim.run_simulation()
     sim.save_results()

jafgen/simulation.py

@@ -24,8 +24,9 @@
 T_8PM = time_from_total_minutes(60 * 20)
 
 class Simulation:
-    def __init__(self, years: int, prefix: str):
+    def __init__(self, years: int, days: int, prefix: str):
         self.years = years
+        self.days = days
         self.scale = 100
         self.prefix = prefix
         self.stores = [
@@ -58,11 +59,11 @@ def __init__(self, years: int, prefix: str):
         self.customers: dict[CustomerId, Customer] = {}
         self.orders: list[Order] = []
         self.tweets: list[Tweet] = []
-        self.sim_days = 365 * self.years
+        self.sim_days = 365 * self.years + self.days
 
     def run_simulation(self):
         for i in track(
-            range(self.sim_days), description="🥪 Pressing fresh jaffles..."
+            range(self.sim_days), description=f"🥪 Pressing {self.sim_days} days of fresh jaffles..."
         ):
             for market in self.markets:
                 day = Day(i)