1 """Fixer that changes plain strings to bytes strings."""

     2 

     3 import re

     4 

     5 from lib2to3 import fixer_base

     6 from lib2to3.pgen2 import token

     7 from lib2to3.fixer_util import Name

     8 from lib2to3.pygram import python_symbols as syms

     9 

    10 _re = re.compile(r'[rR]?[\'\"]')

    11 

    12 # XXX: Implementing a blacklist in 2to3 turned out to be more troublesome than

    13 # blacklisting some modules inside the fixers. So, this is what I came with.

    14 

    15 blacklist = ('mercurial/demandimport.py',

    16              'mercurial/py3kcompat.py', # valid python 3 already

    17              'mercurial/i18n.py',

    18             )

    19 

    20 def isdocstring(node):

    21     def isclassorfunction(ancestor):

    22         symbols = (syms.funcdef, syms.classdef)

    23         # if the current node is a child of a function definition, a class

    24         # definition or a file, then it is a docstring

    25         if ancestor.type == syms.simple_stmt:

    26             try:

    27                 while True:

    28                     if ancestor.type in symbols:

    29                         return True

    30                     ancestor = ancestor.parent

    31             except AttributeError:

    32                 return False

    33         return False

    34 

    35     def ismodule(ancestor):

    36         # Our child is a docstring if we are a simple statement, and our

    37         # ancestor is file_input. In other words, our child is a lone string in

    38         # the source file.

    39         try:

    40             if (ancestor.type == syms.simple_stmt and

    41                 ancestor.parent.type == syms.file_input):

    42                     return True

    43         except AttributeError:

    44             return False

    45 

    46     def isdocassignment(ancestor):

    47         # Assigning to __doc__, definitely a string

    48         try:

    49             while True:

    50                 if (ancestor.type == syms.expr_stmt and

    51                     Name('__doc__') in ancestor.children):

    52                         return True

    53                 ancestor = ancestor.parent

    54         except AttributeError:

    55             return False

    56 

    57     if ismodule(node.parent) or \

    58        isdocassignment(node.parent) or \

    59        isclassorfunction(node.parent):

    60         return True

    61     return False

    62 

    63 def shouldtransform(node):

    64     specialnames = ['__main__']

    65 

    66     if node.value in specialnames:

    67         return False

    68 

    69     ggparent = node.parent.parent.parent

    70     sggparent = str(ggparent)

    71 

    72     if 'getattr' in sggparent or \

    73        'hasattr' in sggparent or \

    74        'setattr' in sggparent or \

    75        'encode' in sggparent or \

    76        'decode' in sggparent:

    77         return False

    78 

    79     return True

    80 

    81 class FixBytes(fixer_base.BaseFix):

    82 

    83     PATTERN = 'STRING'

    84 

    85     def transform(self, node, results):

    86         # The filename may be prefixed with a build directory.

    87         if self.filename.endswith(blacklist):

    88             return

    89         if node.type == token.STRING:

    90             if _re.match(node.value):

    91                 if isdocstring(node):

    92                     return

    93                 if not shouldtransform(node):

    94                     return

    95                 new = node.clone()

    96                 new.value = 'b' + new.value

    97                 return new

    98