2 files changed, 55 insertions, 9 deletions

diff --git a/overlays/worktime/shell.nix b/overlays/worktime/shell.nix
new file mode 100644
index 00000000..18d2a68a
--- /dev/null
+++ b/overlays/worktime/shell.nix
@@ -0,0 +1,5 @@
+{ pkgs ? import <nixpkgs> {} }:
+
+pkgs.mkShell {
+  buildInputs = [(pkgs.python310.withPackages (ps: with ps; [pyxdg python-dateutil uritools requests configparser tabulate]))];
+}
diff --git a/overlays/worktime/worktime.py b/overlays/worktime/worktime.py
index 1fc00061..310fa084 100755
--- a/overlays/worktime/worktime.py
+++ b/overlays/worktime/worktime.py
@@ -27,6 +27,9 @@ from sys import stderr
 
 from tabulate import tabulate
 
+from itertools import groupby
+from functools import cache
+
 class TogglAPISection(Enum):
     TOGGL = '/api/v8'
     REPORTS = '/reports/api/v2'
@@ -118,8 +121,10 @@ class Worktime(object):
     leave_days = set()
     leave_budget = dict()
     time_per_day = None
+    workdays = None
 
     @staticmethod
+    @cache
     def holidays(year):
       holidays = dict()
 
@@ -149,6 +154,13 @@ class Worktime(object):
         config.read(f"{config_dir}/worktime.ini")
         return config
 
+    def ordinal_workday(self, date):
+        start_date = datetime(date.year, 1, 1, tzinfo=tzlocal()).date()
+        return len([1 for offset in range(0, (date - start_date).days + 1) if self.would_be_workday(start_date + timedelta(days = offset))])
+
+    def would_be_workday(self, date):
+        return date.isoweekday() in self.workdays and date not in set(day for (day, val) in Worktime.holidays(date.year).items() if val >= 1)
+
     def __init__(self, start_datetime=None, end_datetime=None, now=None, include_running=True, force_day_to_work=True, **kwargs):
       self.include_running = include_running
       self.force_day_to_work = force_day_to_work
@@ -178,8 +190,8 @@ class Worktime(object):
 
 
       hours_per_week = float(config.get('WORKTIME', 'HoursPerWeek', fallback=40))
-      workdays = set([int(d.strip()) for d in config.get('WORKTIME', 'Workdays', fallback='1,2,3,4,5').split(',')])
-      self.time_per_day = timedelta(hours = hours_per_week) / len(workdays)
+      self.workdays = set([int(d.strip()) for d in config.get('WORKTIME', 'Workdays', fallback='1,2,3,4,5').split(',')])
+      self.time_per_day = timedelta(hours = hours_per_week) / len(self.workdays)
 
       holidays = dict()
 
@@ -237,7 +249,7 @@ class Worktime(object):
                               if end_date.date() < day or day < start_date.date():
                                   continue
 
-                              if excused_kind == 'leave' and not (day in holidays and holidays[day] >= self.time_per_day) and day.isoweekday() in workdays:
+                              if excused_kind == 'leave' and self.would_be_workday(day):
                                   self.leave_days.add(day)
                               holidays[day] = time
           except IOError as e:
@@ -256,7 +268,7 @@ class Worktime(object):
                   if stripped_line:
                       [hours, datestr] = stripped_line.split(' ')
                       constr = datestr.split(',')
-                      for d in [start_day + timedelta(days = x) for x in range(0, (end_day - start_day).days + 1 + int(timedelta(hours = float(hours)).total_seconds() / 60 * (7 / len(workdays)) * 2))]:
+                      for d in [start_day + timedelta(days = x) for x in range(0, (end_day - start_day).days + 1 + int(timedelta(hours = float(hours)).total_seconds() / 60 * (7 / len(self.workdays)) * 2))]:
                           for c in constr:
                               if c in calendar.day_abbr:
                                   if not d.strftime('%a') == c: break
@@ -283,7 +295,7 @@ class Worktime(object):
           end_day = max(end_day, max(list(pull_forward)))
 
       for day in [start_day + timedelta(days = x) for x in range(0, (end_day - start_day).days + 1)]:
-          if day.isoweekday() in workdays:
+          if day.isoweekday() in self.workdays:
               time_to_work = self.time_per_day
               if day in holidays.keys():
                 time_to_work -= holidays[day]
@@ -467,7 +479,8 @@ def holidays(year, **args):
     print(tabulate(table_data, tablefmt="plain"))
 
 def leave(year, table, **args):
-    worktime = Worktime(**dict(**args, end_datetime = datetime(year = year + 1, month = 1, day = 1, tzinfo=tzlocal()) - timedelta(microseconds=1)))
+    def_year = datetime.now(tzlocal()).year
+    worktime = Worktime(**dict(**args, end_datetime = datetime(year = (year if year else def_year) + 1, month = 1, day = 1, tzinfo=tzlocal()) - timedelta(microseconds=1)))
     config = Worktime.config()
     date_format = config.get('WORKTIME', 'DateFormat', fallback='%Y-%m-%d')
     leave_expires = config.get('WORKTIME', 'LeaveExpires', fallback=None)
@@ -475,6 +488,7 @@ def leave(year, table, **args):
         leave_expires = datetime.strptime(leave_expires, '%m-%d').date()
 
     leave_budget = deepcopy(worktime.leave_budget)
+    year_leave_budget = deepcopy(worktime.leave_budget) if year else None
     years = sorted(leave_budget.keys())
     for day in sorted(worktime.leave_days):
         for iyear in years:
@@ -484,18 +498,45 @@ def leave(year, table, **args):
                 continue
 
             leave_budget[iyear] -= 1
+            if year_leave_budget and day.year < year:
+                year_leave_budget[iyear] -= 1
             break
         else:
             print(f'Unaccounted leave: {day}', file=stderr)
 
-    if table:
+    if table and year:
+        table_data = []
+        leave_days = sorted([day for day in worktime.leave_days if day.year == year and worktime.would_be_workday(day)])
+
+        count = 0
+        for _, group in groupby(enumerate(leave_days), lambda kv: kv[0] - worktime.ordinal_workday(kv[1])):
+            group = list(map(lambda kv: kv[1], group))
+
+            for day in group:
+                for iyear in years:
+                    if day > leave_expires.replace(year = iyear + 1):
+                        continue
+                    if year_leave_budget[iyear] <= 0:
+                        continue
+
+                    year_leave_budget[iyear] -= 1
+                    break
+
+            next_count = count + len(group)
+            if len(group) > 1:
+                table_data.append([count, group[0].strftime('%m-%d') + '--' + group[-1].strftime('%m-%d'), len(group), sum(year_leave_budget.values())])
+            else:
+                table_data.append([count, group[0].strftime('%m-%d'), len(group), sum(year_leave_budget.values())])
+            count = next_count
+        print(tabulate(table_data, tablefmt="plain"))
+    elif table:
         table_data = []
         for year, days in leave_budget.items():
             leave_days = sorted([day for day in worktime.leave_days if day.year == year])
             table_data += [[year, days, ','.join(map(lambda d: d.strftime('%m-%d'), leave_days))]]
         print(tabulate(table_data, tablefmt="plain"))
     else:
-        print(leave_budget[year])
+        print(leave_budget[year if year else def_year])
 
 def main():
     parser = argparse.ArgumentParser(prog = "worktime", description = 'Track worktime using toggl API')
@@ -513,7 +554,7 @@ def main():
     holidays_parser.add_argument('year', metavar = 'YEAR', type = int, help = 'Year to evaluate holidays for (default: current year)', default = datetime.now(tzlocal()).year, nargs='?')
     holidays_parser.set_defaults(cmd = holidays)
     leave_parser = subparsers.add_parser('leave')
-    leave_parser.add_argument('year', metavar = 'YEAR', type = int, help = 'Year to evaluate leave days for (default: current year)', default = datetime.now(tzlocal()).year, nargs='?')
+    leave_parser.add_argument('year', metavar = 'YEAR', type = int, help = 'Year to evaluate leave days for (default: current year)', default = None, nargs='?')
     leave_parser.add_argument('--table', action = 'store_true')
     leave_parser.set_defaults(cmd = leave)
     args = parser.parse_args()