squad-track-duration: Improve use of datetime and annotate code

Make code easier to read and understand by switching out use of
strings for datetimes, for making use of the datetime objects.
Annotate the code to clarify what it does.

Signed-off-by: Katie Worton <[email protected]>
Signed-off-by: Anders Roxell <[email protected]>

squad-track-duration

@@ -12,7 +12,7 @@ import json
 import logging
 import os
 import sys
-from datetime import date, timedelta
+from datetime import datetime, timedelta
 from pathlib import Path
 
 import pandas as pd
@@ -45,6 +45,23 @@ class MetaFigure:
         return self.description
 
 
+def parse_datetime_from_string(datetime_string):
+    accepted_datetime_formats = ["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]
+
+    # Loop through each accepted datetime format and try parse it
+    for datetime_format in accepted_datetime_formats:
+        try:
+            # If the format parses successfully, return the datetime object
+            return datetime.strptime(datetime_string, datetime_format)
+        except ValueError:
+            pass
+
+    # If no format can be parsed, raise an argument error
+    raise argparse.ArgumentTypeError(
+        f"Unsupported datetime format {datetime_string}. Accepted formats are {accepted_datetime_formats}"
+    )
+
+
 def parse_args():
     parser = argparse.ArgumentParser(description="Track duration")
 
@@ -62,12 +79,14 @@ def parse_args():
 
     parser.add_argument(
         "--start-datetime",
+        type=parse_datetime_from_string,
         required=True,
         help="Starting date time. Example: 2022-01-01 or 2022-01-01T00:00:00",
     )
 
     parser.add_argument(
         "--end-datetime",
+        type=parse_datetime_from_string,
         required=True,
         help="Ending date time. Example: 2022-12-31 or 2022-12-31T00:00:00",
     )
@@ -108,57 +127,60 @@ def save_build_cache_to_artifactorial(data, days_ago=None):
 
 def get_data(args, build_cache):
     start_datetime = args.start_datetime
-    if "T" not in start_datetime:
-        start_datetime = f"{start_datetime}T00:00:00"
-
     end_datetime = args.end_datetime
-    if "T" not in end_datetime:
-        end_datetime = f"{end_datetime}T23:59:59"
-
 
     group = Squad().group(args.group)
     project = group.project(args.project)
     environments = project.environments(count=ALL).values()
 
-    start_date = start_datetime.split('T')[0]
-    end_date = end_datetime.split('T')[0]
-
-    start_year = int(start_date.split("-")[0])
-    start_month = int(start_date.split("-")[1])
-    start_day = int(start_date.split("-")[2])
-    to_year = int(end_date.split("-")[0])
-    to_month = int(end_date.split("-")[1])
-    to_day = int(end_date.split("-")[2])
-
     first_start_day = True
+    final_end_date = False
     tmp_data = []
 
-    tmp_start_date = date(start_year, start_month, start_day)
-    end_date = date(to_year, to_month, to_day)
+    # Set up a delta which determines how many days of data to read from SQUAD
+    # per loop. Minimum delta is 1 day and delta must be in whole days to keep
+    # this code easy to read, understand and debug.
     delta = timedelta(days=1)
 
-    tmp_start_date -= delta
+    if delta.days < 1:
+        raise Exception("Minimum delta is 1 day for this code to work.")
+    if delta.seconds != 0 or delta.microseconds != 0:
+        raise Exception("Deltas must be whole days only.")
 
-    while tmp_start_date < end_date:
-        tmp_end_date = tmp_start_date + delta
-        tmp_start_date += delta
+    # Loops through each delta until the end date and filters the SQUAD data
+    # for that delta
+    while not final_end_date:
 
+        # If it is the first date in the range, use the provided start datetime
         if first_start_day:
             first_start_day = False
-            start_time = f"T{start_datetime.split('T')[1]}"
+            # Use the provided start time for the first day
+            tmp_start_datetime = start_datetime
         else:
-            start_time = "T00:00:00"
-
-        if tmp_end_date == end_date:
-            to_time = f"T{end_datetime.split('T')[1]}"
+            # For all other days, update the date by the delta then use the
+            # start of the day by zeroing hours, minutes and seconds
+            tmp_start_datetime += delta
+            tmp_start_datetime = tmp_start_datetime.replace(hour=0, minute=0, second=0)
+
+        # If the delta for this iteration sends us over the end of the range,
+        # use the provided end datetime
+        if tmp_start_datetime + delta >= end_datetime:
+            # We have reached the last day, so use this as the end date
+            tmp_end_datetime = end_datetime
+            final_end_date = True
         else:
-            to_time = "T23:59:59"
+            # Otherwise take the start time (with minutes zeroed) + delta
+            tmp_end_datetime = (
+                tmp_start_datetime.replace(hour=0, minute=0, second=0) + delta
+            )
 
-        logger.info(f"Fetching builds from SQUAD, start_datetime: {tmp_start_date}{start_time}, end_datetime: {tmp_end_date}{to_time}")
+        logger.info(
+            f"Fetching builds from SQUAD, start_datetime: {tmp_start_datetime}, end_datetime: {tmp_end_datetime}"
+        )
 
         filters = {
-            "created_at__lt": f"{tmp_end_date}{to_time}",
-            "created_at__gt": f"{tmp_start_date}{start_time}",
+            "created_at__lt": tmp_end_datetime.strftime("%Y-%m-%dT%H:%M:%S"),
+            "created_at__gt": tmp_start_datetime.strftime("%Y-%m-%dT%H:%M:%S"),
             "count": ALL,
         }
 
@@ -259,6 +281,9 @@ def run():
     if args.debug:
         logger.setLevel(level=logging.DEBUG)
 
+    if args.start_datetime > args.end_datetime:
+        raise Exception("Start time must be earlier than end time.")
+
     df = pd.DataFrame(
         {
             "build_name": [],