source: stamper/stamper/stamper.py@ 9:2457f4022c60

import json
import re
from datetime import datetime, timedelta
from os.path import expanduser, exists


STAMPS_FILE = expanduser('~/.workstamps.json')
DATE_FORMAT = '%Y-%m-%d %H:%M:%S'


class Stamper(object):

    def __init__(self, stamps_file=STAMPS_FILE):
        self.stamps_file = STAMPS_FILE
        self.ensure_stamps_file()
        self.stamps = []

    def ensure_stamps_file(self):
        if not exists(self.stamps_file):
            with open(self.stamps_file, 'w') as stamps_file:
                stamps_file.write('')

    def load_stamps(self):
        with open(self.stamps_file, 'r') as stamps_file:
            try:
                self.stamps = json.load(stamps_file)
            except ValueError:
                self.stamps = []

    def save_stamps(self):
        with open(self.stamps_file, 'w') as stamps_file:
            json.dump(self.stamps, stamps_file, indent=4)

    def stamp(self, start, end, customer, action):
        self.stamps.append({
            'start': start,
            'end': end,
            'customer': customer,
            'action': action,
        })

    def last_stamp(self):
        if not self.stamps:
            return None
        return self.stamps[-1]

    def worktime(self, start, end):
        worktime = (datetime.strptime(end, DATE_FORMAT) -
                    datetime.strptime(start, DATE_FORMAT))
        return worktime.seconds

    def validate_filter(self, stamp_filter):
        """
        Validate a given filter. Filters can have the following notation:

        - %Y-%m-%d--%Y-%m-%d: Times recorded between two dates

        - -%Y-%m-%d: Times recorded before a given date

        - +%Y-%m-%d: Times recorded after a given date

        - N...N[d|w|m|y]: Times recorded N...N days/weeks/months/years ago
        """
        # First try the date filters, one by one
        matches = ['%Y-%m-%d', '-%Y-%m-%d', '+%Y-%m-%d']
        for match in matches:
            try:
                if '--' in stamp_filter:
                    filter_from, filter_to = stamp_filter.split('--')
                    filter_from = datetime.strptime(stamp_filter, match)
                    filter_to = datetime.strptime(stamp_filter, match)
                else:
                    valid_filter = datetime.strptime(stamp_filter, match)
            except ValueError:
                pass
            else:
                return stamp_filter

        valid_filter = re.search(r'(\d+[dwmyDWMY]{1})', stamp_filter)
        if valid_filter:
            return stamp_filter

        # Not a valid filter
        return None

    def customers(self):
        customers = []
        for stamp in self.stamps:
            if stamp['customer'] not in customers:
                customers.append(stamp['customer'])
        customers.remove(None)
        return customers

    def totals(self, stamp_filter=None):
        totals = {}
        for stamp in self.stamps:
            customer = stamp['customer']
            if customer:
                # c will be None for "start" stamps, having no end time
                if customer not in totals:
                    totals[customer] = 0
                totals[customer] += self.worktime(stamp['start'], stamp['end'])
        return totals

    def details(self):
        details = {}
        totals = {}
        for stamp in self.stamps:
            if stamp['customer']:
                # avoid "start" stamps
                start_day = stamp['start'].split(' ')[0]
                if start_day not in details:
                    details[start_day] = []
                if start_day not in totals:
                    totals[start_day] = 0
                worktime = self.worktime(stamp['start'], stamp['end'])
                details[start_day].append(
                    ' -> %(worktime)s %(customer)s %(action)s' % {
                        'worktime': str(timedelta(seconds=worktime)),
                        'customer': stamp['customer'],
                        'action': stamp['action']
                    })
                totals[start_day] += worktime
        for day in totals:
            totals[day] = str(timedelta(seconds=totals[day]))
        return details, totals

    def details_by_customer(self, customer):
        details = {}
        totals = {}
        for stamp in self.stamps:
            if stamp['customer'] == customer:
                start_day = stamp['start'].split(' ')[0]
                if start_day not in details:
                    details[start_day] = []
                if start_day not in totals:
                    totals[start_day] = 0
                worktime = self.worktime(stamp['start'], stamp['end'])
                details[start_day].append(
                    ' -> %(worktime)s %(customer)s %(action)s' % {
                        'worktime': str(timedelta(seconds=worktime)),
                        'customer': stamp['customer'],
                        'action': stamp['action']
                    })
                totals[start_day] += worktime
        for day in totals:
            totals[day] = str(timedelta(seconds=totals[day]))
        return details, totals

    def show_stamps(self, customer=None, stamp_filter=None, verbose=False):
        if stamp_filter:
            stamp_filter = self.validate_filter(stamp_filter)

        totals = self.totals(stamp_filter)

        if customer:
            total = timedelta(seconds=totals.get(customer, 0))
            print(' %(customer)s: %(total)s' % {'customer': customer,
                                                'total': total})
        else:
            for c in totals:
                total = timedelta(seconds=totals[c])
                print(' %(customer)s: %(total)s' % {'customer': c,
                                                    'total': total})

        if verbose:
            if customer:
                details, totals = self.details_by_customer(customer)
            else:
                details, totals = self.details()
            for day in details:
                print('------ %(day)s ------' % {'day': day})
                for line in details[day]:
                    print(line)
                print(' Total: %(total)s' % {'total': totals[day]})