worktime: multi-client support

1 files changed, 50 insertions, 7 deletions

diff --git a/overlays/worktime/worktime.py b/overlays/worktime/worktime.py
index 310fa084..166bea1c 100755
--- a/overlays/worktime/worktime.py
+++ b/overlays/worktime/worktime.py
@@ -46,9 +46,10 @@ class TogglAPIError(Exception):
             return self.response.text
 
 class TogglAPI(object):
-    def __init__(self, api_token, workspace_id):
+    def __init__(self, api_token, workspace_id, client_ids):
         self._api_token = api_token
         self._workspace_id = workspace_id
+        self._client_ids = set(map(int, client_ids.split(','))) if client_ids else None
 
     def _make_url(self, api=TogglAPISection.TOGGL, section=['time_entries', 'current'], params={}):
         if api is TogglAPISection.REPORTS:
@@ -87,11 +88,35 @@ class TogglAPI(object):
           elif req_start > start_date:
             req_start = datetime.combine(req_start.astimezone(timezone.utc).date(), time(tzinfo=timezone.utc)) + timedelta(days = 1)
 
-          url = self._make_url(api = TogglAPISection.REPORTS, section = ['summary'], params={'since': req_start.astimezone(timezone.utc).isoformat(), 'until': req_end.astimezone(timezone.utc).isoformat(), 'rounding': rounding})
-          r = self._query(url = url, method='GET')
-          if not r or not r.json():
-              raise TogglAPIError(r)
-          billable_acc += timedelta(milliseconds=r.json()['total_billable']) if r.json()['total_billable'] else timedelta(milliseconds=0)
+          def get_report(client_ids = self._client_ids):
+              nonlocal req_start, req_end, rounding, self
+
+              if client_ids is not None and not client_ids:
+                  return timedelta(milliseconds = 0)
+
+              params = { 'since': req_start.astimezone(timezone.utc).isoformat(),
+                         'until': req_end.astimezone(timezone.utc).isoformat(),
+                         'rounding': rounding,
+                         'billable': 'yes'
+                        }
+              if client_ids is not None:
+                  params |= { 'client_ids': ','.join(map(str, client_ids)) }
+              url = self._make_url(api = TogglAPISection.REPORTS, section = ['summary'], params = params)
+              r = self._query(url = url, method='GET')
+              if not r or not r.json():
+                  raise TogglAPIError(r)
+              res = timedelta(milliseconds=r.json()['total_billable']) if r.json()['total_billable'] else timedelta(milliseconds=0)
+              return res
+
+          if 0 in self._client_ids:
+              url = self._make_url(api = TogglAPISection.TOGGL, section = ['workspaces', self._workspace_id, 'clients'])
+              r = self._query(url = url, method = 'GET')
+              if not r or not r.json():
+                  raise TogglAPIError(r)
+
+              billable_acc += get_report(None) - get_report(set(map(lambda c: c['id'], r.json())))
+
+          billable_acc += get_report(self._client_ids - {0})
 
         return billable_acc
 
@@ -105,6 +130,24 @@ class TogglAPI(object):
         if not r.json()['data'] or not r.json()['data']['billable']:
             return None
 
+        if self._client_ids is not None:
+            if 'pid' in r.json()['data'] and r.json()['data']['pid']:
+                url = self._make_url(api = TogglAPISection.TOGGL, section = ['projects', str(r.json()['data']['pid'])])
+                pr = self._query(url = url, method = 'GET')
+                if not pr or not pr.json():
+                    raise TogglAPIError(pr)
+
+                if not pr.json()['data']:
+                    return None
+
+                if 'cid' in pr.json()['data'] and pr.json()['data']['cid']:
+                    if pr.json()['data']['cid'] not in self._client_ids:
+                        return None
+                elif 0 not in self._client_ids:
+                    return None
+            elif 0 not in self._client_ids:
+                return None
+
         start = isoparse(r.json()['data']['start'])
 
         return now - start if start <= now else None
@@ -170,7 +213,7 @@ class Worktime(object):
 
       config = Worktime.config()
       config_dir = BaseDirectory.load_first_config('worktime')
-      api = TogglAPI(api_token=config['TOGGL']['ApiToken'], workspace_id=config['TOGGL']['Workspace'])
+      api = TogglAPI(api_token=config['TOGGL']['ApiToken'], workspace_id=config['TOGGL']['Workspace'], client_ids=config.get('TOGGL', 'ClientIds', fallback=None))
       date_format = config.get('WORKTIME', 'DateFormat', fallback='%Y-%m-%d')
 
       start_date = start_datetime or datetime.strptime(config['WORKTIME']['StartDate'], date_format).replace(tzinfo=tzlocal())