1 files changed, 17 insertions, 16 deletions

diff --git a/overlays/worktime/worktime/__main__.py b/overlays/worktime/worktime/__main__.py
index 21316182..c2c1829c 100755
--- a/overlays/worktime/worktime/__main__.py
+++ b/overlays/worktime/worktime/__main__.py
@@ -6,6 +6,8 @@ from xdg import BaseDirectory
 import toml
 from uritools import (uricompose)
 
+from inspect import signature
+
 from dateutil.easter import *
 from dateutil.tz import *
 from dateutil.parser import isoparse
@@ -633,14 +635,15 @@ def holidays(year, table_format, **args):
 
 def leave(year, table, table_format, **args):
     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", {}).get("DateFormat", '%Y-%m-%d')
     leave_expires = config.get("WORKTIME", {}).get("LeaveExpires", None)
     if leave_expires:
         leave_expires = datetime.strptime(leave_expires, '%m-%d').date()
+    worktime = Worktime(**dict(**args, end_datetime = datetime.combine(leave_expires.replace(year = (year if year else def_year) + 1), time(), tzinfo=tzlocal()) + timedelta(days=1) if leave_expires else datetime(year = (year if year else def_year) + 1, month = 1, day = 1, tzinfo=tzlocal()) - timedelta(microseconds=1)))
+    date_format = config.get("WORKTIME", {}).get("DateFormat", '%Y-%m-%d')
 
     days = [worktime.start_date.date() + timedelta(days = x) for x in range(0, (worktime.end_date.date() - worktime.start_date.date()).days + 1)]
+    leave_days_budget = {}
 
     leave_budget = deepcopy(worktime.leave_budget)
     year_leave_budget = deepcopy(worktime.leave_budget) if year else None
@@ -655,36 +658,28 @@ def leave(year, table, table_format, **args):
             if leave_budget[iyear] <= 0:
                 continue
 
+            leave_days_budget[day] = iyear
             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 and year:
         table_data = []
-        leave_days = sorted([day for day in worktime.leave_days if day.year == year and worktime.would_be_workday(day)])
+        leave_days = sorted([day for day in worktime.leave_days if leave_days_budget[day] == 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
+                year_leave_budget[year] -= 1
 
             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())])
+                table_data.append([count, group[0].strftime('%m-%d' if group[0].year == year else '%Y-%m-%d') + '…' + group[-1].strftime('%m-%d' if group[-1].year == year else '%Y-%m-%d'), len(group), year_leave_budget[year]])
             else:
-                table_data.append([count, group[0].strftime('%m–%d'), len(group), sum(year_leave_budget.values())])
+                table_data.append([count, group[0].strftime('%m-%d' if group[0].year == year else '%Y-%m-%d'), len(group), year_leave_budget[year]])
             count = next_count
         print(tabulate(table_data, tablefmt=table_format, headers=["Running Count", "Leave Days", "# of Leave Days", "# of Leave Days Left"] if table_format != 'plain' else None))
     elif table:
@@ -834,7 +829,13 @@ def main():
         classification_parser.set_defaults(cmd = partial(classification, classification_name=classification_name))
     args = parser.parse_args()
 
-    args.cmd(**vars(args))
+    args.cmd(
+        **{
+            k: v
+            for k, v in vars(args).items()
+            if k in signature(args.cmd).parameters.keys()
+        }
+    )
 
 if __name__ == "__main__":
     sys.exit(main())