1 files changed, 97 insertions, 17 deletions

diff --git a/overlays/worktime/worktime.py b/overlays/worktime/worktime.py
index 678c0d69..9cfc6cd4 100755
--- a/overlays/worktime/worktime.py
+++ b/overlays/worktime/worktime.py
@@ -14,7 +14,7 @@ from dateutil.parser import isoparse
 
 from enum import Enum
 
-from math import (copysign, ceil)
+from math import (copysign, ceil, floor)
 
 import calendar
 
@@ -23,6 +23,7 @@ import argparse
 from copy import deepcopy
 
 import sys
+from sys import stderr
 
 from tabulate import tabulate
 
@@ -114,6 +115,8 @@ class Worktime(object):
     include_running = True
     time_to_work = None
     force_day_to_work = True
+    leave_days = set()
+    leave_budget = dict()
 
     @staticmethod
     def holidays(year):
@@ -179,25 +182,67 @@ class Worktime(object):
 
       holidays = dict()
 
+      leave_per_year = int(config.get('WORKTIME', 'LeavePerYear', fallback=30))
       for year in range(start_date.year, end_date.year + 1):
           holidays |= {k: v * time_per_day for k, v in Worktime.holidays(year).items()}
+          leave_frac = 1
+          if date(year, 1, 1) < start_date.date():
+              leave_frac = (date(year + 1, 1, 1) - start_date.date()) / (date(year + 1, 1, 1) - date(year, 1, 1))
+          self.leave_budget |= {year: floor(leave_per_year * leave_frac)}
 
       try:
-          with open(f"{config_dir}/excused", 'r') as excused:
+          with open(f"{config_dir}/reset-leave", 'r') as excused:
               for line in excused:
                   stripped_line = line.strip()
                   if stripped_line:
-                      splitLine = stripped_line.split(' ')
-                      if len(splitLine) == 2:
-                          [hours, date] = splitLine
-                          day = datetime.strptime(date, date_format).replace(tzinfo=tzlocal()).date()
-                          holidays[day] = timedelta(hours = float(hours))
-                      else:
-                          holidays[datetime.strptime(stripped_line, date_format).replace(tzinfo=tzlocal()).date()] = time_per_day
+                      [datestr, count] = stripped_line.split(' ')
+                      day = datetime.strptime(datestr, date_format).replace(tzinfo=tzlocal()).date()
+                      if day != start_date.date():
+                          continue
+                      
+                      self.leave_budget[day.year] = (self.leave_budget[day.year] if day.year in self.leave_budget else 0) + int(count)
       except IOError as e:
           if e.errno != 2:
               raise e
 
+
+      for excused_kind in {'excused', 'leave'}:
+          try:
+              with open(f"{config_dir}/{excused_kind}", 'r') as excused:
+                  for line in excused:
+                      stripped_line = line.strip()
+                      if stripped_line:
+                          splitLine = stripped_line.split(' ')
+                          fromDay = toDay = None
+                          def parse_datestr(datestr):
+                              nonlocal fromDay, toDay
+                              def parse_single(singlestr):
+                                  return datetime.strptime(singlestr, date_format).replace(tzinfo=tzlocal()).date()
+                              if '--' in datestr:
+                                  [fromDay,toDay] = datestr.split('--')
+                                  fromDay = parse_single(fromDay)
+                                  toDay = parse_single(toDay)
+                              else:
+                                  fromDay = toDay = parse_single(datestr)
+                          time = time_per_day
+                          if len(splitLine) == 2:
+                              [hours, datestr] = splitLine
+                              time = timedelta(hours = float(hours))
+                              parse_datestr(datestr)
+                          else:
+                              parse_datestr(stripped_line)
+
+                          for day in [fromDay + timedelta(days = x) for x in range(0, (toDay - fromDay).days + 1)]:
+                              if end_date.date() < day or day < start_date.date():
+                                  continue
+
+                              if excused_kind == 'leave' and not (day in holidays and holidays[day] >= time_per_day) and day.isoweekday() in workdays:
+                                  self.leave_days.add(day)
+                              holidays[day] = time
+          except IOError as e:
+              if e.errno != 2:
+                  raise e
+
       pull_forward = dict()
       
       start_day = start_date.date()
@@ -208,8 +253,8 @@ class Worktime(object):
               for line in excused:
                   stripped_line = line.strip()
                   if stripped_line:
-                      [hours, date] = stripped_line.split(' ')
-                      constr = date.split(',')
+                      [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 c in constr:
                               if c in calendar.day_abbr:
@@ -253,8 +298,8 @@ class Worktime(object):
                   if stripped_line:
                       splitLine = stripped_line.split(' ')
                       if len(splitLine) == 2:
-                          [hours, date] = splitLine
-                          day = datetime.strptime(date, date_format).replace(tzinfo=tzlocal()).date()
+                          [hours, datestr] = splitLine
+                          day = datetime.strptime(datestr, date_format).replace(tzinfo=tzlocal()).date()
                           extra_days_to_work[day] = timedelta(hours = float(hours))
                       else:
                           extra_days_to_work[datetime.strptime(stripped_line, date_format).replace(tzinfo=tzlocal()).date()] = time_per_day
@@ -395,19 +440,49 @@ def diff(now, **args):
 
     print(now.time_to_work - then.time_to_work)
 
-def holidays(now, **args):
+def holidays(year, **args):
     config = Worktime.config()
     date_format = config.get('WORKTIME', 'DateFormat', fallback='%Y-%m-%d')
 
     table_data = []
     
-    holidays = Worktime.holidays(now.year)
+    holidays = Worktime.holidays(year)
     for k, v in holidays.items():
         kstr = k.strftime(date_format)
-        
         table_data += [[kstr, v]]
     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)))
+    config = Worktime.config()
+    date_format = config.get('WORKTIME', 'DateFormat', fallback='%Y-%m-%d')
+    leave_expires = config.get('WORKTIME', 'LeaveExpires', fallback=None)
+    if leave_expires:
+        leave_expires = datetime.strptime(leave_expires, '%m-%d').date()
+
+    leave_budget = deepcopy(worktime.leave_budget)
+    years = sorted(leave_budget.keys())
+    for day in sorted(worktime.leave_days):
+        for iyear in years:
+            if day > leave_expires.replace(year = iyear + 1):
+                continue
+            if leave_budget[iyear] <= 0:
+                continue
+
+            leave_budget[iyear] -= 1
+            break
+        else:
+            print(f'Unaccounted leave: {day}', file=stderr)
+    
+    if 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])
+
 def main():
     parser = argparse.ArgumentParser(prog = "worktime", description = 'Track worktime using toggl API')
     parser.add_argument('--time', dest = 'now', metavar = 'TIME', type = lambda s: datetime.fromisoformat(s).replace(tzinfo=tzlocal()), help = 'Time to calculate status for (default: current time)', default = datetime.now(tzlocal()))
@@ -421,7 +496,12 @@ def main():
     diff_parser = subparsers.add_parser('diff')
     diff_parser.set_defaults(cmd = diff)
     holidays_parser = subparsers.add_parser('holidays')
+    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('--table', action = 'store_true')
+    leave_parser.set_defaults(cmd = leave)
     args = parser.parse_args()
 
     args.cmd(**vars(args))