#!/usr/bin/env python3

#

# byteify-strings.py - transform string literals to be Python 3 safe

#

# Copyright 2015 Gregory Szorc <gregory.szorc@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

import argparse

import contextlib

import errno

import os

import sys

import tempfile

import token

import tokenize

def adjusttokenpos(t, ofs):

    """Adjust start/end column of the given token"""

    return t._replace(

        start=(t.start[0], t.start[1] + ofs), end=(t.end[0], t.end[1] + ofs)

    )

def replacetokens(tokens, opts):

    """Transform a stream of tokens from raw to Python 3.

    Returns a generator of possibly rewritten tokens.

    The input token list may be mutated as part of processing. However,

    its changes do not necessarily match the output token stream.

    """

    sysstrtokens = set()

    # The following utility functions access the tokens list and i index of

    # the for i, t enumerate(tokens) loop below

    def _isop(j, *o):

        """Assert that tokens[j] is an OP with one of the given values"""

        try:

            return tokens[j].type == token.OP and tokens[j].string in o

        except IndexError:

            return False

    def _findargnofcall(n):

        """Find arg n of a call expression (start at 0)

        Returns index of the first token of that argument, or None if

        there is not that many arguments.

        Assumes that token[i + 1] is '('.

        """

        nested = 0

        for j in range(i + 2, len(tokens)):

            if _isop(j, ')', ']', '}'):

                # end of call, tuple, subscription or dict / set

                nested -= 1

                if nested < 0:

                    return None

            elif n == 0:

                # this is the starting position of arg

                return j

            elif _isop(j, '(', '[', '{'):

                nested += 1

            elif _isop(j, ',') and nested == 0:

                n -= 1

        return None

    def _ensuresysstr(j):

        """Make sure the token at j is a system string

        Remember the given token so the string transformer won't add

        the byte prefix.

        Ignores tokens that are not strings. Assumes bounds checking has

        already been done.

        """

        k = j

        currtoken = tokens[k]

        while currtoken.type in (token.STRING, token.NEWLINE, tokenize.NL):

            k += 1

            if currtoken.type == token.STRING and currtoken.string.startswith(

                ("'", '"')

            ):

                sysstrtokens.add(currtoken)

            try:

                currtoken = tokens[k]

            except IndexError:

                break

    def _isitemaccess(j):

        """Assert the next tokens form an item access on `tokens[j]` and that

        `tokens[j]` is a name.

        """

        try:

            return (

                tokens[j].type == token.NAME

                and _isop(j + 1, '[')

                and tokens[j + 2].type == token.STRING

                and _isop(j + 3, ']')

            )

        except IndexError:

            return False

    def _ismethodcall(j, *methodnames):

        """Assert the next tokens form a call to `methodname` with a string

        as first argument on `tokens[j]` and that `tokens[j]` is a name.

        """

        try:

            return (

                tokens[j].type == token.NAME

                and _isop(j + 1, '.')

                and tokens[j + 2].type == token.NAME

                and tokens[j + 2].string in methodnames

                and _isop(j + 3, '(')

                and tokens[j + 4].type == token.STRING

            )

        except IndexError:

            return False

    coldelta = 0  # column increment for new opening parens

    coloffset = -1  # column offset for the current line (-1: TBD)

    parens = [(0, 0, 0, -1)]  # stack of (line, end-column, column-offset, type)

    ignorenextline = False  # don't transform the next line

    insideignoreblock = False  # don't transform until turned off

    for i, t in enumerate(tokens):

        # Compute the column offset for the current line, such that

        # the current line will be aligned to the last opening paren

        # as before.

        if coloffset < 0:

            lastparen = parens[-1]

            if t.start[1] == lastparen[1]:

                coloffset = lastparen[2]

            elif t.start[1] + 1 == lastparen[1] and lastparen[3] not in (

                token.NEWLINE,

                tokenize.NL,

            ):

                # fix misaligned indent of s/util.Abort/error.Abort/

                coloffset = lastparen[2] + (lastparen[1] - t.start[1])

            else:

                coloffset = 0

        # Reset per-line attributes at EOL.

        if t.type in (token.NEWLINE, tokenize.NL):

            yield adjusttokenpos(t, coloffset)

            coldelta = 0

            coloffset = -1

            if not insideignoreblock:

                ignorenextline = (

                    tokens[i - 1].type == token.COMMENT

                    and tokens[i - 1].string == "# no-py3-transform"

                )

            continue

        if t.type == token.COMMENT:

            if t.string == "# py3-transform: off":

                insideignoreblock = True

            if t.string == "# py3-transform: on":

                insideignoreblock = False

        if ignorenextline or insideignoreblock:

            yield adjusttokenpos(t, coloffset)

            continue

        # Remember the last paren position.

        if _isop(i, '(', '[', '{'):

            parens.append(t.end + (coloffset + coldelta, tokens[i + 1].type))

        elif _isop(i, ')', ']', '}'):

            parens.pop()

        # Convert most string literals to byte literals. String literals

        # in Python 2 are bytes. String literals in Python 3 are unicode.

        # Most strings in Mercurial are bytes and unicode strings are rare.

        # Rather than rewrite all string literals to use ``b''`` to indicate

        # byte strings, we apply this token transformer to insert the ``b``

        # prefix nearly everywhere.

        if t.type == token.STRING and t not in sysstrtokens:

            s = t.string

            # Preserve docstrings as string literals. This is inconsistent

            # with regular unprefixed strings. However, the

            # "from __future__" parsing (which allows a module docstring to

            # exist before it) doesn't properly handle the docstring if it

            # is b''' prefixed, leading to a SyntaxError. We leave all

            # docstrings as unprefixed to avoid this. This means Mercurial

            # components touching docstrings need to handle unicode,

            # unfortunately.

            if s[0:3] in ("'''", '"""'):

                # If it's assigned to something, it's not a docstring

                if not _isop(i - 1, '='):

                    yield adjusttokenpos(t, coloffset)

                    continue

            # If the first character isn't a quote, it is likely a string

            # prefixing character (such as 'b', 'u', or 'r'. Ignore.

            if s[0] not in ("'", '"'):

                yield adjusttokenpos(t, coloffset)

                continue

            # String literal. Prefix to make a b'' string.

            yield adjusttokenpos(t._replace(string='b%s' % t.string), coloffset)

            coldelta += 1

            continue

        # This looks like a function call.

        if t.type == token.NAME and _isop(i + 1, '('):

            fn = t.string

            # *attr() builtins don't accept byte strings to 2nd argument.

            if (

                fn

                in (

                    'getattr',

                    'setattr',

                    'hasattr',

                    'safehasattr',

                    'wrapfunction',

                    'wrapclass',

                    'addattr',

                )

                and (opts['allow-attr-methods'] or not _isop(i - 1, '.'))

            ):

                arg1idx = _findargnofcall(1)

                if arg1idx is not None:

                    _ensuresysstr(arg1idx)

            # .encode() and .decode() on str/bytes/unicode don't accept

            # byte strings on Python 3.

            elif fn in ('encode', 'decode') and _isop(i - 1, '.'):

                for argn in range(2):

                    argidx = _findargnofcall(argn)

                    if argidx is not None:

                        _ensuresysstr(argidx)

            # It changes iteritems/values to items/values as they are not

            # present in Python 3 world.

            elif opts['dictiter'] and fn in ('iteritems', 'itervalues'):

                yield adjusttokenpos(t._replace(string=fn[4:]), coloffset)

                continue

        if t.type == token.NAME and t.string in opts['treat-as-kwargs']:

            if _isitemaccess(i):

                _ensuresysstr(i + 2)

            if _ismethodcall(i, 'get', 'pop', 'setdefault', 'popitem'):

                _ensuresysstr(i + 4)

        # Looks like "if __name__ == '__main__'".

        if (

            t.type == token.NAME

            and t.string == '__name__'

            and _isop(i + 1, '==')

        ):

            _ensuresysstr(i + 2)

        # Emit unmodified token.

        yield adjusttokenpos(t, coloffset)

def process(fin, fout, opts):

    tokens = tokenize.tokenize(fin.readline)

    tokens = replacetokens(list(tokens), opts)

    fout.write(tokenize.untokenize(tokens))

def tryunlink(fname):

    try:

        os.unlink(fname)

    except OSError as err:

        if err.errno != errno.ENOENT:

            raise

@contextlib.contextmanager

def editinplace(fname):

    n = os.path.basename(fname)

    d = os.path.dirname(fname)

    fp = tempfile.NamedTemporaryFile(

        prefix='.%s-' % n, suffix='~', dir=d, delete=False

    )

    try:

        yield fp

        fp.close()

        if os.name == 'nt':

            tryunlink(fname)

        os.rename(fp.name, fname)

    finally:

        fp.close()

        tryunlink(fp.name)

def main():

    ap = argparse.ArgumentParser()

    ap.add_argument(

        '--version', action='version', version='Byteify strings 1.0'

    )

    ap.add_argument(

        '-i',

        '--inplace',

        action='store_true',

        default=False,

        help='edit files in place',

    )

    ap.add_argument(

        '--dictiter',

        action='store_true',

        default=False,

        help='rewrite iteritems() and itervalues()',

    ),

    ap.add_argument(

        '--allow-attr-methods',

        action='store_true',

        default=False,

        help='also handle attr*() when they are methods',

    ),

    ap.add_argument(

        '--treat-as-kwargs',

        nargs="+",

        default=[],

        help="ignore kwargs-like objects",

    ),

    ap.add_argument('files', metavar='FILE', nargs='+', help='source file')

    args = ap.parse_args()

    opts = {

        'dictiter': args.dictiter,

        'treat-as-kwargs': set(args.treat_as_kwargs),

        'allow-attr-methods': args.allow_attr_methods,

    }

    for fname in args.files:

        fname = os.path.realpath(fname)

        if args.inplace:

            with editinplace(fname) as fout:

                with open(fname, 'rb') as fin:

                    process(fin, fout, opts)

        else:

            with open(fname, 'rb') as fin:

                fout = sys.stdout.buffer

                process(fin, fout, opts)

if __name__ == '__main__':

    if sys.version_info[0:2] < (3, 7):

        print('This script must be run under Python 3.7+')

        sys.exit(3)

    main()