From 6580e10a33675e305faa94d56b9949f5ff485256 Mon Sep 17 00:00:00 2001
From: Pedram Bakh <56321501+PedramBakh@users.noreply.github.com>
Date: Mon, 11 Sep 2023 16:42:13 +0200
Subject: [PATCH] Add API support for CLI tool

---
 .github/workflows/publish.yml |  4 ++--
 carbontracker/cli.py          | 10 +++++++---
 carbontracker/tracker.py      |  6 ++++--
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 7ce3559..19b8a89 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -52,7 +52,7 @@ jobs:
       run: flake8 carbontracker --count --select=E9,F63,F7,F82 --show-source --statistics
 
     - name: Format with Black
-      run: black --check --line-length 120 carbontracker
+      run: black --line-length 120 carbontracker
 
   build-n-publish:
     needs: [test, lint]
@@ -71,7 +71,7 @@ jobs:
 
       - name: Build the package
         run: python -m build
-        
+
       - name: Publish to Test PyPI
         if: startsWith(github.ref, 'refs/tags/dev-v')
         uses: pypa/gh-action-pypi-publish@v1.4.2
diff --git a/carbontracker/cli.py b/carbontracker/cli.py
index d70dc6e..27f5aac 100644
--- a/carbontracker/cli.py
+++ b/carbontracker/cli.py
@@ -1,17 +1,21 @@
 import argparse
 import subprocess
 from carbontracker.tracker import CarbonTracker
+import ast
 
 
 def main():
     parser = argparse.ArgumentParser(description="CarbonTracker CLI")
     parser.add_argument("--script", type=str, required=True, help="Python file to execute")
     parser.add_argument("--log_dir", type=str, help="Log directory", default="./logs")
-    # Add more arguments as needed
-
+    parser.add_argument("--api_keys", type=str, help="API keys in a dictionary-like format, e.g., "
+                                                     "'{\"electricitymaps\": \"YOUR_KEY\"}'", default=None)
     args = parser.parse_args()
 
-    tracker = CarbonTracker(epochs=1, log_dir=args.log_dir, epochs_before_pred=0)
+    # Parse the API keys string into a dictionary
+    api_keys = ast.literal_eval(args.api_keys) if args.api_keys else None
+
+    tracker = CarbonTracker(epochs=1, log_dir=args.log_dir, epochs_before_pred=0, api_keys=api_keys)
     tracker.epoch_start()
 
     # Execute script with python
diff --git a/carbontracker/tracker.py b/carbontracker/tracker.py
index f86d861..2aed5f1 100644
--- a/carbontracker/tracker.py
+++ b/carbontracker/tracker.py
@@ -367,8 +367,10 @@ def _output_actual(self):
         time = np.sum(times)
         _co2eq = self._co2eq(energy)
         conversions = co2eq.convert(_co2eq) if self.interpretable else None
-
-        self._output_energy(f"Actual consumption for {self.epoch_counter} epoch(s):", time, energy, _co2eq, conversions)
+        if self.epochs_before_pred == 0:
+            self._output_energy("Actual consumption:", time, energy, _co2eq, conversions)
+        else:
+            self._output_energy(f"Actual consumption for {self.epoch_counter} epoch(s):", time, energy, _co2eq, conversions)
 
     def _output_pred(self):
         """Output predicted usage for full training epochs."""