"""
Enable optimizations on the string module.
"""
from astoptimizer.compatibility import (
    PYTHON2, PYTHON3,
    NATIVE_STR, INT_TYPES, STR_TYPES)
from astoptimizer.config import Function
import string

def setup_config(config):
    CONSTANTS = (
        'ascii_letters', 'ascii_lowercase', 'ascii_uppercase',
        'digits', 'hexdigits', 'octdigits',
        'printable', 'punctuation', 'whitespace')
    if PYTHON2:
        CONSTANTS += ('letters', 'lowercase', 'uppercase')
    for name in CONSTANTS:
        config.add_constant('string.' + name, getattr(string, name))

    if PYTHON3:
        # on Python2, capwords() is locale dependent
        config.add_func('string.capwords', Function(string.capwords, 1, NATIVE_STR))
    if PYTHON2:
        config.add_func('string.atof', Function(string.atof, 1, STR_TYPES, catch=ValueError))
        config.add_func('string.atoi', Function(string.atoi, (1, 2), STR_TYPES, INT_TYPES, catch=ValueError))
        config.add_func('string.atol', Function(string.atol, (1, 2), STR_TYPES, INT_TYPES, catch=ValueError))
        # locale dependent?
        # capitalize, capwords
        config.add_func('string.expandtabs', Function(string.expandtabs, (1, 2), STR_TYPES, INT_TYPES))
        # TODO: center, count, find, index, join, joinfields, ljust, lower,
        # TODO: lstrip, maketrans, replace, rfind, rindex, rjust, rsplit,
        # TODO: rstrip, split, splitfields, strip, swapcase, translate, upper,
        # TODO: zfill

