1# -*- coding: utf-8 -*-
2"""
3Part of the astor library for Python AST manipulation.
4
5License: 3-clause BSD
6
7Copyright (c) 2015 Patrick Maupin
8
9Pretty-print strings for the decompiler
10
11We either return the repr() of the string,
12or try to format it as a triple-quoted string.
13
14This is a lot harder than you would think.
15
16This has lots of Python 2 / Python 3 ugliness.
17
18"""
19
20import re
21
22try:
23    special_unicode = unicode
24except NameError:
25    class special_unicode(object):
26        pass
27
28try:
29    basestring = basestring
30except NameError:
31    basestring = str
32
33
34def _properly_indented(s, line_indent):
35    mylist = s.split('\n')[1:]
36    mylist = [x.rstrip() for x in mylist]
37    mylist = [x for x in mylist if x]
38    if not s:
39        return False
40    counts = [(len(x) - len(x.lstrip())) for x in mylist]
41    return counts and min(counts) >= line_indent
42
43
44mysplit = re.compile(r'(\\|\"\"\"|\"$)').split
45replacements = {'\\': '\\\\', '"""': '""\\"', '"': '\\"'}
46
47
48def _prep_triple_quotes(s, mysplit=mysplit, replacements=replacements):
49    """ Split the string up and force-feed some replacements
50        to make sure it will round-trip OK
51    """
52
53    s = mysplit(s)
54    s[1::2] = (replacements[x] for x in s[1::2])
55    return ''.join(s)
56
57
58def string_triplequote_repr(s):
59    """Return string's python representation in triple quotes.
60    """
61    return '"""%s"""' % _prep_triple_quotes(s)
62
63
64def pretty_string(s, embedded, current_line, uni_lit=False,
65                  min_trip_str=20, max_line=100):
66    """There are a lot of reasons why we might not want to or
67       be able to return a triple-quoted string.  We can always
68       punt back to the default normal string.
69    """
70
71    default = repr(s)
72
73    # Punt on abnormal strings
74    if (isinstance(s, special_unicode) or not isinstance(s, basestring)):
75        return default
76    if uni_lit and isinstance(s, bytes):
77        return 'b' + default
78
79    len_s = len(default)
80
81    if current_line.strip():
82        len_current = len(current_line)
83        second_line_start = s.find('\n') + 1
84        if embedded > 1 and not second_line_start:
85            return default
86
87        if len_s < min_trip_str:
88            return default
89
90        line_indent = len_current - len(current_line.lstrip())
91
92        # Could be on a line by itself...
93        if embedded and not second_line_start:
94            return default
95
96        total_len = len_current + len_s
97        if total_len < max_line and not _properly_indented(s, line_indent):
98            return default
99
100    fancy = string_triplequote_repr(s)
101
102    # Sometimes this doesn't work.  One reason is that
103    # the AST has no understanding of whether \r\n was
104    # entered that way in the string or was a cr/lf in the
105    # file.  So we punt just so we can round-trip properly.
106
107    try:
108        if eval(fancy) == s and '\r' not in fancy:
109            return fancy
110    except Exception:
111        pass
112    return default
113