1# -*- coding: utf-8 -*-
2
3"""
4clint.colored
5~~~~~~~~~~~~~
6
7This module provides a simple and elegant wrapper for colorama.
8
9"""
10
11import os
12import re
13import sys
14
15import colorama
16
17PY3 = sys.version_info[0] >= 3
18
19__all__ = (
20    'red', 'green', 'yellow', 'blue',
21    'black', 'magenta', 'cyan', 'white', 'normal',
22    'clean', 'disable',
23)
24
25COLORS = __all__[:-2]
26
27if 'get_ipython' in dir():
28    """
29       when ipython is fired lot of variables like _oh, etc are used.
30       There are so many ways to find current python interpreter is ipython.
31       get_ipython is easiest is most appealing for readers to understand.
32    """
33    DISABLE_COLOR = True
34else:
35    DISABLE_COLOR = False
36
37
38class ColoredString(object):
39    """Enhanced string for __len__ operations on Colored output."""
40    def __init__(self, color, s, always_color=False, bold=False):
41        super(ColoredString, self).__init__()
42        if not PY3 and isinstance(s, unicode):
43            self.s = s.encode('utf-8')
44        else:
45            self.s = s
46        self.color = color
47        self.always_color = always_color
48        self.bold = bold
49        if os.environ.get('PIPENV_FORCE_COLOR'):
50            self.always_color = True
51
52    def __getattr__(self, att):
53        def func_help(*args, **kwargs):
54            result = getattr(self.s, att)(*args, **kwargs)
55            try:
56                is_result_string = isinstance(result, basestring)
57            except NameError:
58                is_result_string = isinstance(result, str)
59            if is_result_string:
60                return self._new(result)
61            elif isinstance(result, list):
62                return [self._new(x) for x in result]
63            else:
64                return result
65        return func_help
66
67    @property
68    def color_str(self):
69        style = 'BRIGHT' if self.bold else 'NORMAL'
70        c = '%s%s%s%s%s' % (getattr(colorama.Fore, self.color), getattr(colorama.Style, style), self.s, colorama.Fore.RESET, getattr(colorama.Style, 'NORMAL'))
71
72        if self.always_color:
73            return c
74        elif sys.stdout.isatty() and not DISABLE_COLOR:
75            return c
76        else:
77            return self.s
78
79    def __len__(self):
80        return len(self.s)
81
82    def __repr__(self):
83        return "<%s-string: '%s'>" % (self.color, self.s)
84
85    def __unicode__(self):
86        value = self.color_str
87        if isinstance(value, bytes):
88            return value.decode('utf8')
89        return value
90
91    if PY3:
92        __str__ = __unicode__
93    else:
94        def __str__(self):
95            return self.color_str
96
97    def __iter__(self):
98        return iter(self.color_str)
99
100    def __add__(self, other):
101        return str(self.color_str) + str(other)
102
103    def __radd__(self, other):
104        return str(other) + str(self.color_str)
105
106    def __mul__(self, other):
107        return (self.color_str * other)
108
109    def _new(self, s):
110        return ColoredString(self.color, s)
111
112
113def clean(s):
114    strip = re.compile("([^-_a-zA-Z0-9!@#%&=,/'\";:~`\$\^\*\(\)\+\[\]\.\{\}\|\?\<\>\\]+|[^\s]+)")
115    txt = strip.sub('', str(s))
116
117    strip = re.compile(r'\[\d+m')
118    txt = strip.sub('', txt)
119
120    return txt
121
122
123def normal(string, always=False, bold=False):
124    return ColoredString('RESET', string, always_color=always, bold=bold)
125
126
127def black(string, always=False, bold=False):
128    return ColoredString('BLACK', string, always_color=always, bold=bold)
129
130
131def red(string, always=False, bold=False):
132    return ColoredString('RED', string, always_color=always, bold=bold)
133
134
135def green(string, always=False, bold=False):
136    return ColoredString('GREEN', string, always_color=always, bold=bold)
137
138
139def yellow(string, always=False, bold=False):
140    return ColoredString('YELLOW', string, always_color=always, bold=bold)
141
142
143def blue(string, always=False, bold=False):
144    return ColoredString('BLUE', string, always_color=always, bold=bold)
145
146
147def magenta(string, always=False, bold=False):
148    return ColoredString('MAGENTA', string, always_color=always, bold=bold)
149
150
151def cyan(string, always=False, bold=False):
152    return ColoredString('CYAN', string, always_color=always, bold=bold)
153
154
155def white(string, always=False, bold=False):
156    # This upsets people...
157    return ColoredString('WHITE', string, always_color=always, bold=bold)
158
159
160def disable():
161    """Disables colors."""
162    global DISABLE_COLOR
163
164    DISABLE_COLOR = True
165