contrib/hgfixes/fix_bytes.py
changeset 28400 a84fc98a8af5
parent 28399 2af0156cebaa
child 28401 2565fe39a76c
equal deleted inserted replaced
28399:2af0156cebaa 28400:a84fc98a8af5
     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