timezones

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os, sys, argparse, datetime, dateutil.zoneinfo, dateutil.relativedelta
from extras import ColoredArgParser
from extras import PrettyTable
from extras import ColoredSetting
from extras import has_fuzzy_matched
from extras import has_substr_matched

WIDTH_ARGUMENTS = 20
WIDTH_NAMES = 40
WIDTH_LOCATIONS = 40
WIDTH_CODES = 10
WIDTH_UTCOFFSETS = 20

def parse_args():
    parser = ColoredArgParser(description = 'List or find the time zones.')
    parser.add_argument('-i', '--infile',
        action = 'append',
        type = argparse.FileType(mode = 'r', encoding = 'UTF-8'),
        help = 'Input file.')
    parser.add_argument('-l', '--list',
        default = False,
        action = 'store_true',
        help = 'List all timezone.')
    parser.add_argument('-f', '--fuzzy',
        action = 'store_true',
        default = False,
        dest = 'is_fuzzy',
        help = 'Enable fuzzy matching.')
    parser.add_argument('--color',
        choices = ['auto', 'always', 'never'],
        default = 'auto',
        dest = 'colored',
        help = 'When to colored, the default is auto.')
    parser.add_argument('words',
        nargs = '*',
        help = 'Search for timezone contains all the given words.')
    return parser.parse_args()

def format_timedelta_to_utcoffset(td):
    rd = dateutil.relativedelta.relativedelta(seconds = int(td.total_seconds()))
    return 'UTC{}{:d}{}'.format('-' if td.total_seconds() < 0 else '+', abs(rd.hours), (':' + str(abs(rd.minutes)) if rd.minutes else ''))

def listall():
    dt = datetime.datetime(2015,1,1)
    pt = PrettyTable(enable_painting = ColoredSetting().is_colorize(sys.stdout))
    pt.add_field('', alignment = '<', fixed_width = WIDTH_ARGUMENTS)
    pt.add_field('Name', alignment = '<', fixed_width = WIDTH_NAMES)
    pt.add_field('Location', alignment = '<', fixed_width = WIDTH_LOCATIONS)
    pt.add_field('Code', alignment = '<', fixed_width = WIDTH_CODES)
    pt.add_field('UTC offset', alignment = '<', fixed_width = WIDTH_UTCOFFSETS)
    for k, v in dateutil.zoneinfo.get_zonefile_instance().zones.items():
        pt.add_record('', k, v, v.tzname(dt), format_timedelta_to_utcoffset(v.utcoffset(dt)))
    print(pt)

def listfound(words, is_fuzzy):
    if not words:
        return
    dt = datetime.datetime(2015,1,1)
    pt = PrettyTable(enable_painting = ColoredSetting().is_colorize(sys.stdout))
    pt.add_field('Argument', alignment = '>', fixed_width = WIDTH_ARGUMENTS)
    pt.add_field('Name', alignment = '<', fixed_width = WIDTH_NAMES)
    pt.add_field('Location', alignment = '<', fixed_width = WIDTH_LOCATIONS)
    pt.add_field('Code', alignment = '<', fixed_width = WIDTH_CODES)
    pt.add_field('UTC offset', alignment = '<', fixed_width = WIDTH_UTCOFFSETS)
    for word in words:
        for name, v in dateutil.zoneinfo.get_zonefile_instance().zones.items():
            location = str(v)
            code = str(v.tzname(dt))
            utcoffset = format_timedelta_to_utcoffset(v.utcoffset(dt))
            haystack = (name, location, code, utcoffset)
            if has_substr_matched(word, haystack, ignore_case = True)   \
                or (is_fuzzy and has_fuzzy_matched(word, haystack, ignore_case = True)):
                pt.add_record(repr(word), name, location, code, utcoffset)
    print(pt)

def main(args = None):
    if args is None:
        return 0
    ColoredSetting(args.colored)
    if args.list:
        listall()
    words = args.words
    if not args.list and not args.words and not args.infile:
        args.infile = ( sys.stdin, )
    for file in args.infile or ():
        if file.name == '<stdin>':
            print('Press Ctrl-D when finished.', flush = True)
        words.extend(file.read().split())
    listfound(words, args.is_fuzzy)

if __name__ == '__main__':
    sys.exit(main(parse_args()))