• Jérome Perrin's avatar
    *: rewrite with lib2to3.fixes.fix_asserts and ad-hoc assertin · 2e366054
    Jérome Perrin authored
    The add-hoc assertin filter:
    
    --
    
    from typing import List
    
    import lib2to3
    
    from lib2to3.fixer_base import BaseFix
    from lib2to3.fixer_util import Comma, Name
    
    class FixAssertIn(BaseFix):
    
      PATTERN = """
          power< any+ trailer< '.' meth=("assertTrue" | "assertFalse")>
          trailer< '('
            comparison< (needle=any ( comp_op<'not' 'in'> | 'in' ) haystack=any) >
          ')' > >
      """
    
      def transform(self, node: lib2to3.pytree.Node,
                    results: List[lib2to3.pytree.Base]):
    
        needle = results['needle']
        haystack = results['haystack']
        meth = results["meth"][0]
    
        method_map = {True: 'assertIn', False: 'assertNotIn'}
        method_in = meth.value == 'assertTrue'
        if 'not' in str(needle.parent.children[1]):
          method_in = not method_in
        meth.replace(Name(method_map[method_in], prefix=meth.prefix))
    
        needle.parent.children = [needle, Comma(), haystack]
    2e366054