Option to only download, pass desired pathes and dark theme

main.py

@@ -9,12 +9,12 @@
 SHELLY_BASE_URL = "http://192.168.178.99"
 PHASE_URL = SHELLY_BASE_URL + "/emeter/{}/em_data.csv"  # Replace {} with 0, 1, 2
 
-CSV_PATH = Path(__file__).parent / "data/input"
+DATA_PATH = Path(__file__).parent / "data"
 
 logger = logging.getLogger(__name__)
 
 
-def download_data():
+def download_data(csv_path: Path):
     phases = [
         pd.read_csv(PHASE_URL.format(i), parse_dates=["Date/time UTC"])
         for i in tqdm(range(3))
@@ -26,14 +26,14 @@ def download_data():
 
     # Split to one dataframe per day
     df["Date"] = df["Date/time UTC"].dt.date
-    (CSV_PATH / "current_day.csv").unlink(missing_ok=True)
+    (csv_path / "current_day.csv").unlink(missing_ok=True)
     for date in tqdm(df["Date"].unique(), desc="Saving data"):
         day_df = df[df["Date"] == date]
         if date == pd.Timestamp.now(tz="UTC").date():
-            file_path = CSV_PATH / f"current_day.csv"
+            file_path = csv_path / f"current_day.csv"
         else:
             year, month = date.year, date.month
-            file_path = CSV_PATH / f"{year}/{month:02}/{date}.csv"
+            file_path = csv_path / f"{year}/{month:02}/{date}.csv"
             if file_path.exists():
                 tqdm.write(f"Skipping {date} as it already exists")
                 continue
@@ -47,16 +47,15 @@ def download_data():
 
 
 def main(
-        download=False,
+        csv_path: Path,
+        output_path: Path,
         sample_rate: str | None = None,
         plot_phases=False,
         timeframe_start: str | None = None,
         timeframe_end: str | None = None,
         y_axis_limit: int | None = None,
+        dark_theme=False,
 ):
-    if download:
-        download_data()
-
     if timeframe_start:
         timeframe_start = pd.Timestamp(timeframe_start)
     if timeframe_end:
@@ -66,7 +65,7 @@ def main(
         raise ValueError("timeframe_start must be before timeframe_end")
 
     # Read the data
-    files = list(CSV_PATH.rglob("*.csv"))
+    files = list(csv_path.rglob("*.csv"))
     dfs = []
     for file in sorted(files):
         if (timeframe_start or timeframe_end) and file.stem != "current_day":
@@ -75,7 +74,6 @@ def main(
                 continue
             if timeframe_end and timeframe_end.date() < file_date:
                 continue
-        print(f"Reading {file}")
         df = pd.read_csv(file, parse_dates=["Date/time UTC"])
         dfs.append(df)
     df = pd.concat(dfs)
@@ -98,6 +96,21 @@ def main(
         ts = ts.resample(sample_rate).sum()
 
     # Plot the data
+    if dark_theme:
+        plt.style.use("dark_background")
+        # colors of the graph should look normal
+        plt.rcParams["axes.prop_cycle"] = plt.cycler(color=[
+            "#1f77b4",
+            "#ff7f0e",
+            "#2ca02c",
+            "#d62728",
+            "#9467bd",
+            "#8c564b",
+            "#e377c2",
+            "#7f7f7f",
+            "#bcbd22",
+            "#17becf"
+        ])
     if plot_phases:
         ax = ts.plot.area(
             y=["Active energy Wh (A)", "Active energy Wh (B)", "Active energy Wh (C)"],
@@ -119,31 +132,68 @@ def main(
             ylim=(0, y_axis_limit) if y_axis_limit else None,
         )
     else:
-        ts.plot(y=["Total Active energy Wh", "Total Returned energy Wh"],
-                title="Total Active and Returned energy Wh",
-                xlabel="Date/time UTC",
-                ylabel="Energy Wh",
-                figsize=(20, 10),
-                grid=True,
-                style=["-", ":", "--", "-.", "-", ":", "--", "-."],
-                drawstyle="steps-post",
-                ylim=(0, y_axis_limit) if y_axis_limit else None,
-                )
+        ts.plot(
+            y=["Total Active energy Wh", "Total Returned energy Wh"],
+            title="Total Active and Returned energy Wh",
+            xlabel="Date/time UTC",
+            ylabel="Energy Wh",
+            figsize=(20, 10),
+            grid=True,
+            style=["-", ":", "--", "-.", "-", ":", "--", "-."],
+            drawstyle="steps-post",
+            ylim=(0, y_axis_limit) if y_axis_limit else None,
+        )
 
     # Save the plot
-    plt.savefig("data/plot.png")
+    plt.savefig(output_path)
 
     print(f"Total Active energy Wh: {ts['Total Active energy Wh'].sum()}")
     print(f"Total Returned energy Wh: {ts['Total Returned energy Wh'].sum()}")
 
 
 if __name__ == "__main__":
-    args = argparse.ArgumentParser()
+    args = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    args.add_argument(
+        "--csv-path",
+        type=Path,
+        help="Directory to read the CSV files from. "
+             "If combined with --download, the downloaded files will be saved here.",
+        default=DATA_PATH / "input",
+    )
+    args.add_argument(
+        "--output-path",
+        type=Path,
+        help="Path to save the plot to. "
+             "File ending determines the format. "
+             "Format must be supported by matplotlib."
+             ".png and .svg at least are supported.",
+        default=DATA_PATH / "plot.png",
+    )
     args.add_argument("--download", action="store_true", help="Download the data from the Shelly device")
+    args.add_argument(
+        "--only-download",
+        action="store_true",
+        help="Do not plot, only download the data from the Shelly device"
+    )
     args.add_argument("--sample-rate", type=str, help="Sample rate for the data")
     args.add_argument("--plot-phases", action="store_true", help="Plot the data for each phase")
     args.add_argument("--timeframe-start", type=str, help="Plot the graph starting at this datetime")
     args.add_argument("--timeframe-end", type=str, help="Plot the graph ending at this datetime")
     args.add_argument("--y-axis-limit", type=int, help="Limit the y-axis to this value")
+    args.add_argument("--dark-theme", action="store_true", help="Use a dark theme for the plot")
     args = args.parse_args()
-    main(args.download, args.sample_rate, args.plot_phases, args.timeframe_start, args.timeframe_end, args.y_axis_limit)
+    if args.download or args.only_download:
+        download_data(args.csv_path)
+    if not args.only_download:
+        main(
+            args.csv_path,
+            args.output_path,
+            args.sample_rate,
+            args.plot_phases,
+            args.timeframe_start,
+            args.timeframe_end,
+            args.y_axis_limit,
+            args.dark_theme,
+        )