import ast
import math
import os
import re
import sys
import unittest
import warnings

from astoptimizer import Config as _Config, parse_ast, compile_ast
from astoptimizer.compatibility import (
    PYTHON2, PYTHON27, PYTHON3, PYTHON34,
    u, b,
    BYTES_TYPE, UNICODE_TYPE)
from astoptimizer.optimizer import Namespace, Optimizer

class TestCase(unittest.TestCase):
    if not hasattr(unittest.TestCase, 'skipTest'):
        def skipTest(self, reason):
            print("skip test: %s" % reason)
            return reason

    if not hasattr(unittest.TestCase, 'assertIn'):
        def assertIn(self, item, data):
            self.assertTrue(item in data, "%r not in %r" % (item, data))


class BaseTestCase(TestCase):
    TEXT_PASS = 'Module(body=[Pass()])'

    def __init__(self, *args):
        TestCase.__init__(self, *args)
        self.maxDiff = 4096
        self.infos = []
        self.warnings = []

    def setUp(self):
        del self.warnings[:]
        self._default_config = self.create_default_config()

    def tearDown(self):
        self._check_infos(None)
        self._check_warnings(None)

    def log_info(self, message):
        self.infos.append(message)

    def log_warning(self, message):
        self.warnings.append(message)

    def create_config(self, *features):
        config = _Config(*features)
        config.info_func = self.log_info
        config.warning_func = self.log_warning
        return config

    def create_default_config(self):
        return self.create_config()

    def _text_module(self, expr):
        return 'Module(body=[Expr(value=%s)])' % expr

    def _text_num(self, value):
        return 'Num(n=%r)' % value

    def text_num(self, value):
        return self._text_module(self._text_num(value))

    def _text_name(self, value):
        return 'Name(id=%r, ctx=Load())' % value

    if PYTHON34:
        def _text_name_constant(self, value):
            return 'NameConstant(value=%s)' % value

    def text_name(self, value):
        return self._text_module(self._text_name(value))

    def _text_bool(self, value):
        name = "True" if value else "False"
        if PYTHON34:
            return self._text_name_constant(name)
        else:
            return self._text_name(name)

    def text_bool(self, value):
        return self._text_module(self._text_bool(value))

    if PYTHON3:
        def _text_bytes(self, value):
            return "Bytes(s=%r)" % value
        def _text_unicode(self, value):
            return "Str(s=%r)" % value
    else:
        def _text_bytes(self, value):
            assert isinstance(value, BYTES_TYPE)
            return "Str(s=%r)" % value
        def _text_unicode(self, value):
            assert isinstance(value, UNICODE_TYPE)
            return "Str(s=%r)" % value

    def text_bytes(self, value):
        return self._text_module(self._text_bytes(value))

    def text_unicode(self, value):
        return self._text_module(self._text_unicode(value))

    def _text_native_str(self, value):
        return "Str(s=%r)" % value

    def text_native_str(self, value):
        return self._text_module(self._text_native_str(value))

    def _text_item(self, value):
        if value is None:
            if PYTHON34:
                return self._text_name_constant(None)
            else:
                return self._text_name('None')
        elif isinstance(value, (int, float, complex)):
            return self._text_num(value)
        elif isinstance(value, BYTES_TYPE):
            return self._text_bytes(value)
        elif isinstance(value, UNICODE_TYPE):
            return self._text_unicode(value)
        else:
            raise NotImplementedError("Unknown type: %s" % type(value))

    def _text_tuple(self, *elts):
        elts = (self._text_item(elt) for elt in elts)
        elts = '[%s]' % ', '.join(elts)
        return 'Tuple(elts=%s, ctx=Load())' % elts

    def text_tuple(self, *elts):
        return self._text_module(self._text_tuple(*elts))

    def _text_call(self, name, *args):
        func = self._text_name(name)
        args = list(args)
        return (
            'Call(func=%s, args=[%s], '
            'keywords=[], starargs=None, kwargs=None)'
            % (func, ', '.join(args)))

    def _text_set(self, *elts):
        if PYTHON27:
            elts = (self._text_item(elt) for elt in elts)
            elts = '[%s]' % ', '.join(elts)
            return 'Set(elts=%s)' % elts
        else:
            if elts:
                arg = self._text_tuple(*elts)
                return self._text_call('set', arg)
            else:
                return self._text_call('set')

    def text_set(self, *elts):
        return self._text_module(self._text_set(*elts))

    def text_ast(self, code):
        tree = parse_ast(code)
        return ast.dump(tree)

    def _optimize_ast(self, tree, config=None, catch_syntaxerror=False):
        if config is None:
            config = self._default_config
        optimizer = Optimizer(config)
        tree = optimizer.optimize(tree, None)
        # Ensure that the tree is compilable to bytecode
        try:
            compile_ast(tree)
        except SyntaxError:
            if not catch_syntaxerror:
                raise
        return tree, optimizer

    def _check_messages(self, what, messages, message):
        emitted = messages[:]
        del messages[:]
        if message:
            if isinstance(message, (list, tuple)):
                expected = message
            else:
                expected = [message]
            self.assertEqual(len(emitted), len(expected), emitted)
            for emitted, message in zip(emitted, expected):
                if isinstance(message, str):
                    self.assertEqual(emitted, message)
                else:
                    # regex
                    self.assertTrue(message.search(emitted), emitted)
        elif emitted:
            raise Exception("%s: %s" % (what, emitted))

    def _check_infos(self, message):
        self._check_messages("INFOS", self.infos, message)

    def _check_warnings(self, message):
        self._check_messages("WARNINGS", self.warnings, message)

    def check(self, code, expected, config=None, info=None, warning=None):
        tree = parse_ast(code)
        tree, optimizer = self._optimize_ast(tree, config)
        text = ast.dump(tree)
        self.assertEqual(text, expected)
        self._check_infos(info)
        self._check_warnings(warning)

    def check_not_optimized(self, code,
                            config=None, catch_syntaxerror=False,
                            info=None, warning=None):
        old_tree = parse_ast(code)
        old = ast.dump(old_tree)
        old = re.sub(r"UnaryOp\(op=USub\(\), operand=Num\(n=([0-9]+(?:\.[0-9]*)?)\)\)", lambda regs: "Num(n=-%s)" % regs.group(1), old)
        new_tree, optimizer = self._optimize_ast(old_tree, config, catch_syntaxerror=catch_syntaxerror)
        new = ast.dump(new_tree)
        self.assertEqual(new, old)
        self._check_infos(info)
        self._check_warnings(warning)

    def check_qualnames(self, code, expected,
                        info=None, warning=None):
        tree = parse_ast(code)
        tree, optimizer = self._optimize_ast(tree)
        self.assertEqual(optimizer.namespace.qualnames, expected)
        self._check_infos(info)
        self._check_warnings(warning)


class TestOptimizer(BaseTestCase):
    def test_None(self):
        self.check('(None,) * 3', self.text_tuple(None, None, None))

        config = self.create_config()
        config.add_constant('x', None)
        self.check('x', self.text_ast('None'), config)

    def test_BinOp_int(self):
        self.check('1 + 2', self.text_num(3))
        self.check('10 - 1', self.text_num(9))
        self.check('2 * 3', self.text_num(6))
        self.check('21 // 5', self.text_num(4))
        self.check_not_optimized('21 // 0')
        self.check('21 % 5', self.text_num(1))
        self.check_not_optimized('21 % 0')
        self.check('10 ** 3', self.text_num(1000))
        self.check('-25 ** 0', self.text_num(-1))
        self.check('-2 ** -1', self.text_num(-0.5))
        self.check('1 << 5', self.text_num(32))
        self.check('32 >> 3', self.text_num(4))
        self.check_not_optimized('3 << -1')
        self.check('5 & 3', self.text_num(1))
        self.check('5 | 3', self.text_num(7))
        self.check('5 ^ 3', self.text_num(6))
        if PYTHON3:
            self.check('21 / 5', self.text_num(4.2))
        else:
            self.check_not_optimized('21 / 5')
        self.check_not_optimized('21 / 0')

        # complex operations
        self.check('1 + 2 * 3', self.text_num(7))

        # simple optimizations
        self.check_not_optimized('0 + x')
        self.check_not_optimized('0.0 + x')
        self.check_not_optimized('x + 0')
        self.check_not_optimized('x + 0.0')
        self.check_not_optimized('1 * x')
        self.check_not_optimized('1.0 * x')
        self.check_not_optimized('x * 1')
        self.check_not_optimized('x * 1.0')

        # test min_int / max_int limits
        config = self.create_config()
        config.min_int = -10
        config.max_int = 10
        self.check('1 + 2', self.text_num(3), config)
        self.check_not_optimized('4 * 5', config)
        self.check_not_optimized('2 - 40', config)

    def test_BinOp_float(self):
        self.check('1.0 + 2', self.text_num(3.0))
        self.check('10 - 1.0', self.text_num(9.0))
        self.check('2.0 * 3.0', self.text_num(6.0))
        self.check('21.0 // 5.0', self.text_num(4.0))
        self.check('21.0 % 5.0', self.text_num(1.0))
        self.check('10.0 ** 3', self.text_num(1000.0))
        self.check_not_optimized('(-25) ** 0.5')

        self.check_not_optimized('1.0 << 5')
        self.check_not_optimized('32.0 >> 3')
        self.check_not_optimized('5.0 & 3')
        self.check_not_optimized('5.0 | 3')
        self.check_not_optimized('5.0 ^ 3')

    def test_BinOp_complex(self):
        self.check('1j + 2', self.text_num(1j + 2.0))
        self.check('10 - 1j', self.text_num(10.0 - 1j))
        self.check('2j * 3.0', self.text_num(6j))
        self.check('21.0 / 5.0', self.text_num(4.2))
        if PYTHON3:
            self.check_not_optimized('21j // 5j')
            self.check_not_optimized('21j % 5j')
        else:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", DeprecationWarning)
                self.check('21j // 5j', self.text_num(4+0j))
                self.check('21j % 5j', self.text_num(1j))
        self.check_not_optimized('10j ** 3')
        self.check_not_optimized('2 ** 1j')
        self.check_not_optimized('1.0 << 5')
        self.check_not_optimized('32.0 >> 3')
        self.check_not_optimized('5.0 & 3')
        self.check_not_optimized('5.0 | 3')
        self.check_not_optimized('5.0 ^ 3')

    def test_BinOp_numbers(self):
        self.check('1 + 2.0', self.text_num(3.0))
        self.check('1.0 + 2', self.text_num(3.0))
        self.check('1.0 + 2 * 3.0', self.text_num(7.0))

    def test_BinOp_string(self):
        self.check('%r * 3' % b("a"), self.text_bytes(b('aaa')))
        self.check('%r * 3' % (u('a'),), self.text_unicode(u('aaa')))
        self.check('%r %% 123' % u("x=%s"), self.text_unicode(u('x=123')))
        self.check('%r %% (2, 3)' % u("x=%s y=%s"), self.text_unicode(u('x=2 y=3')))
        if PYTHON2:
            self.check('"x=%s" % 123',
                       self.text_bytes(b('x=123')))
            self.check('"x=%s y=%s" % (2, 3)',
                       self.text_bytes(b('x=2 y=3')))
            self.check('"x\xe9=%s" % 4',
                       self.text_bytes(b('x\xe9=4')))
            self.check('"x=%s" % u"abc"', self.text_unicode(u('x=abc')))
            self.check_not_optimized('"x\xe9=%s" % u"abc"')
        else:
            self.check_not_optimized('b"x=%s" % 123')
            self.check_not_optimized('b"x=%s y=%s" % (2, 3)')
            self.check_not_optimized('"x=%s" % b"abc"')
        self.check_not_optimized('"a" - "b"')
        self.check_not_optimized('"a" * "b"')
        self.check_not_optimized('"a" * 2.0')
        self.check_not_optimized('%r %% x' % b("x=%s"))
        self.check_not_optimized('%r %% 2' % b("x=%s y=%s"))
        self.check_not_optimized('%r %% x' % u("x=%s"))
        self.check_not_optimized('%r %% 2' % u("x=%s y=%s"))

        config = self.create_config()
        config.max_string_length = 5
        self.check('"a" * 5', self.text_native_str('aaaaa'), config)
        self.check_not_optimized('"a" * 6', config)

    def test_BinOp_tuple(self):
        self.check('(1,) + (2,)', self.text_tuple(1, 2))
        self.check('(1, 2) * 2', self.text_tuple(1, 2, 1, 2))
        self.check('2 * (1, 2)', self.text_tuple(1, 2, 1, 2))

        config = self.create_config()
        config.max_tuple_length = 3
        self.check('(1,) * 3', self.text_tuple(1, 1, 1), config)
        self.check_not_optimized('(1,) * 4', config)

    def test_BoolOp(self):
        # a and b constant
        self.check('True and True', self.text_bool(True))
        self.check('False or True', self.text_bool(True))
        self.check('13 or 14', self.text_num(13))

        # a constant
        self.check('True and x', self.text_name('x'))
        self.check('False and x', self.text_bool(False))
        self.check('True or x', self.text_bool(True))
        self.check('False or x', self.text_name('x'))

        # b constant
        self.check_not_optimized('x and True')
        self.check_not_optimized('x and False')
        self.check_not_optimized('x or True')
        self.check_not_optimized('x or False')

        # others
        self.check('4 and 5 and x and 6', self.text_ast('x and 6'))
        self.check('4 and 5 and ~6', self.text_num(-7))

    def test_UnaryOp(self):
        self.check('not True', self.text_bool(False))
        self.check('not 3', self.text_bool(False))
        self.check('not False', self.text_bool(True))
        self.check('not 0', self.text_bool(True))
        self.check_not_optimized('not x')

        self.check('~3', self.text_num(-4))
        self.check_not_optimized('~x')
        self.check_not_optimized('~1.0')

        self.check('-(0.0)', self.text_num(-0.0))
        self.check('-(2)', self.text_num(-2))
        self.check('-(-(3))', self.text_num(3))
        self.check('+(0.0)', self.text_num(0.0))
        self.check('+(2)', self.text_num(2))
        self.check('+(-(3))', self.text_num(-3))
        self.check_not_optimized('-"abc"')
        self.check_not_optimized('+(1, 2)')

        config = self.create_config()
        config.min_int = -10
        config.max_int = 10
        self.check('~9', self.text_num(-10), config)
        self.check_not_optimized('~10', config)

    def test_Subscript(self):
        self.check('%r[:3]' % b("abcdef"), self.text_bytes(b("abc")))
        if PYTHON3:
            self.check('%r[1]' % b("abcdef"), self.text_num(ord(b("b"))))
        else:
            self.check('%r[1]' % b("abcdef"), self.text_bytes(b("b")))
        self.check('"abcdef"[:3]', self.text_native_str("abc"))
        self.check('"abcdef"[2:4]', self.text_native_str("cd"))
        self.check('"abcdef"[2::-1]', self.text_native_str("cba"))
        self.check('"abcdef"[2:0:-1]', self.text_native_str("cb"))
        self.check('"abcdef"[1]', self.text_native_str("b"))
        if sys.version_info < (3,3):
            self.check_not_optimized('%r[0]' % u("\U0010ffff"))
            self.check_not_optimized('%r[2]' % u("abc\U0010ffff-"))
            self.check_not_optimized('%r[4]' % u("abc\U0010ffff-"))
            self.check_not_optimized('%r[-1]' % u("abc\U0010ffff-"))
            self.check_not_optimized('%r[::-1]' % u("abc\U0010ffff-"))
            self.check_not_optimized('%r[3::-1]' % u("abc\U0010ffff-"))
            self.check_not_optimized('%r[3:0:-1]' % u("abc\U0010ffff-"))

            config = self.create_config("pythonbin")
            self.check('%r[0]' % u("\U0010ffff"), self.text_unicode(u('\U0010ffff')[0]), config)
            self.check('%r[2]' % u("abc\U0010ffff-"), self.text_unicode(u('c')), config)
        else:
            self.check('%r[3]' % u("abc\U0010ffff-"), self.text_unicode(u("\U0010ffff")))
            self.check('%r[-1]' % u("abc\U0010ffff-"), self.text_unicode(u("-")))

        self.check('"abc"[2]', self.text_native_str('c'))
        self.check_not_optimized('"abc"[3]')
        self.check('"abc"[-3]', self.text_native_str('a'))
        self.check('"abc"[None:]', self.text_native_str('abc'))
        self.check_not_optimized('"abc"[-4]')
        self.check_not_optimized('"abc"[None]')
        self.check_not_optimized('"abc"[1.0]')
        self.check_not_optimized('"abc"[:1.0]')
        self.check_not_optimized('1[1]')
        self.check_not_optimized('1[:1]')

    def check_pass(self, code, config=None):
        expected = code.replace("pass; pass", "pass")
        expected = expected.replace("\n pass\n pass", "\n pass")
        self.check(code, self.text_ast(expected), config=config)

    def test_Pass(self):
        self.check_not_optimized('pass')
        self.check_pass('pass; pass')
        self.check('pass; func(); pass', self.text_ast('func()'))
        self.check('if True:\n pass\n pass\nif False: pass', self.TEXT_PASS)

        self.check_pass('class Klass:\n pass\n pass')
        self.check_pass('def f():\n pass\n pass')
        config = self.create_config()
        config.unroll_limit = 0
        self.check('for i in (1, 2, 3):\n pass\n print(x)\n pass',
                   self.text_ast('for i in (1, 2, 3):\n print(x)'),
                   config=config)
        self.check_pass('while x:\n pass\n pass')

        code = '\n'.join((
            'try:',
            '    pass; func(); pass',
            'except:',
            '    pass; func(); pass',
            'else:',
            '    pass; func(); pass',
            'finally:',
            '    pass; func(); pass',
        ))
        code2 = '\n'.join((
            'try:',
            '    func()',
            'except:',
            '    func()',
            'else:',
            '    func()',
            'finally:',
            '    func()',
        ))
        self.check(code, self.text_ast(code2))

    def test_If(self):
        self.check('if 0: print("log")', self.TEXT_PASS)
        self.check('if 1: 3', self.text_num(3))
        self.check('if 2**2**2 - 16: 3', self.TEXT_PASS)
        self.check('if 0: pass\nelif 0: pass\npass', self.TEXT_PASS)
        self.check('def f():\n if 4 < 5:\n  return "yes"\n else:\n  return "no"',
                   self.text_ast('def f():\n return "yes"'))
        self.check('def f():\n yield 3\n if 0:\n  yield',
                   self.text_ast('def f():\n yield 3'))
        self.check('def f():\n if g():\n  if True:\n   pass',
                   self.text_ast('def f():\n if g():\n  pass'))
        self.check('if a:\n if b:\n  pass',
                   self.text_ast('if a and b:\n pass'))
        self.check_not_optimized('def f():\n if 0:\n  yield')
        self.check_not_optimized('def f():\n if 1:\n  pass\n else:\n  yield')
        self.check_not_optimized('def f():\n if 0:\n  yield\n yield 3')
        self.check_not_optimized('def f():\n yield 3\n if 0:\n  return 3', catch_syntaxerror=True)
        code = """
def f():
    if test:
        if 1:
            x = 1
            return x
        """.strip()
        expected = """
def f():
    if test:
        x = 1
        return x
        """.strip()
        self.check(code, self.text_ast(expected))

    def test_remove_dead_code(self):
        # return, raise
        self.check('def f():\n return 1\n return 2',
                   self.text_ast('def f():\n return 1'),
                   info='<string>:3: Remove unreachable code: Return(value=Num(n=2))')
        self.check('def f():\n raise\n return 2',
                   self.text_ast('def f():\n raise'),
                   info='<string>:3: Remove unreachable code: Return(value=Num(n=2))')
        self.check('def f():\n g()\n return 1\n h()\n return 2',
                   self.text_ast('def f():\n g()\n return 1'),
                   info=["<string>:4: Remove unreachable code: Expr(value=Call(func=Name(id='h', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None))",
                        "<string>:5: Remove unreachable code: Return(value=Num(n=2))"])
        self.check('def f():\n g()\n raise ValueError("error")\n h()\n return 2',
                   self.text_ast('def f():\n g()\n raise ValueError("error")'),
                   info=["<string>:4: Remove unreachable code: Expr(value=Call(func=Name(id='h', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None))",
                         "<string>:5: Remove unreachable code: Return(value=Num(n=2))"])

        # global, nonlocal
        self.check_not_optimized('''
def f():
    if 0:
        global x
    x = 2
        '''.strip())
        if PYTHON3:
            self.check_not_optimized('''
def f():
    x = 1
    def g():
        if 0:
            nonlocal x
        x = 2
    return g
            '''.strip())

        # generators
        self.check_not_optimized('def f():\n return\n yield 1')
        self.check_not_optimized('def f():\n return 22\n yield 1', catch_syntaxerror=True)
        code = """
def f():
    if 0:
        lambda x:  x        # shouldn't trigger here
        return              # or here
        def f(i):
            return 2*i      # or here
        if 0:
            return 3        # but *this* sucks (line 8)
    if 0:
        yield 2             # because it's a generator (line 10)
        """.strip()
        self.check_not_optimized(code, catch_syntaxerror=True)

        # empty try/except
        self.check('try: pass\nexcept: func1()',
                   self.text_ast('pass'),
                   info=re.compile('Remove dead code: Try'))
        self.check('try: pass\nexcept: func1()\nelse: func2()',
                   self.text_ast('func2()'),
                   info=re.compile('Remove dead code: Try'))
        self.check('try: pass\nfinally: func1()',
                   self.text_ast('func1()'),
                   info=re.compile('Remove dead code: Try'))
        if sys.version_info >= (3, 3):
            infos = re.compile('Remove dead code: Try')
        else:
            infos = (re.compile('Remove dead code: Try'),
                     re.compile('Remove dead code: Try'))
        self.check('try: pass\nexcept: func()\nfinally: func2()',
                   self.text_ast('func2()'),
                   info=infos)
        code = 'try: pass\nexcept: func1()\nelse: func2()\nfinally: func3()'
        self.check(code,
                   self.text_ast('try: func2()\nfinally: func3()'),
                   info=re.compile('Remove dead code: Try'))

        # For
#       TODO: remove empty loop
        #self.check('for x in (1, 2, 3): pass',
        #           self.text_ast('pass'),
        #           info=re.compile('Remove dead code: For'))

    def test_remove_docstring(self):
        keep_docstring = self.create_config()
        keep_docstring.remove_docstring = False
        self.check('def f():\n "docstring"\n return 1',
                   self.text_ast('def f():\n "docstring"\n return 1'),
                   config=keep_docstring)

        no_docstring = self.create_config()
        no_docstring.remove_docstring = True
        self.check('def f():\n "docstring"\n return 1',
                   self.text_ast('def f():\n return 1'),
                   config=no_docstring)

    def test_remove_almost_dead_code(self):
        config = self.create_config()
        config.remove_almost_dead_code = True
        self.check('obj.attr',
                   self.text_ast('pass'),
                   config=config,
                   info=re.compile('Remove dead code'))
        self.check_not_optimized('func()', config=config)

    def test_dont_remove_dead_code(self):
        # disable removal of dead code
        config = self.create_config()
        config.remove_dead_code = False
        config.unroll_limit = 0
        self.check_not_optimized('def f():\n return 1\n return 2', config)
        self.check('if False: print("log")',
                   self.text_ast('if 0: print("log")'),
                   config)
        self.check('while "": print("log")',
                   self.text_ast('while 0: print("log")'),
                   config)
        self.check('2 if "abc" else 3',
                   self.text_ast('2 if 1 else 3'),
                   config)
        self.check('2 if "" else 3',
                   self.text_ast('2 if 0 else 3'),
                   config)
        self.check_not_optimized('for x in (): pass', config)
        self.check_not_optimized('for x in (1, 2, 3): pass', config)
        self.check_not_optimized('[x for x in () if 1]', config)
        self.check('(x for x in "" if True)',
                   self.text_ast('(x for x in "" if 1)'),
                   config)
        self.check('[x for x in "abc" if False]',
                   self.text_ast('[x for x in "abc" if 0]'),
                   config)
        self.check_not_optimized('(x for x in "abc" if 0)', config)

        if PYTHON27:
            self.check('{x*2 for x in (1, 2, 3) if False}',
                       self.text_ast('{x*2 for x in (1, 2, 3) if 0}'),
                       config)
            self.check('{x: x*2 for x in (1, 2, 3) if False}',
                       self.text_ast('{x: x*2 for x in (1, 2, 3) if 0}'),
                       config)

    def test_IfExp(self):
        self.check('4 if "abc" else 5', self.text_num(4))
        self.check('4 if 0 else 5', self.text_num(5))
        self.check_not_optimized('4 if x else 5')

    def test_While(self):
        self.check('while 0: print("log")', self.TEXT_PASS)
        self.check_not_optimized('def f():\n while 0:\n  yield')
        self.check('while "abc": pass',
                   self.text_ast('while 1: pass'))
        self.check('while True: pass',
                   self.text_ast('while 1: pass'))
        self.check('def f():\n yield 3\n while 0:\n  yield 5',
                   self.text_ast('def f():\n yield 3'))
        self.check_not_optimized('while 1: print("log")')
        self.check_not_optimized('def f():\n while 0:\n  yield 5')
        self.check_not_optimized('def f():\n while 0:\n  yield 5\n yield 3')
        self.check_not_optimized('i=0\nwhile i < 5: i += 1')

        code = '\n'.join((
            'while 0:',
            '    print("log")',
            'else:',
            '    print("abc")',
        ))
        self.check(code, self.text_ast('print("abc")'))

    def test_Compare(self):
        self.check('"abc" == "abc"', self.text_bool(True))
        self.check('1 != 2', self.text_bool(True))
        self.check('2 < 2', self.text_bool(False))
        self.check('2 <= 2', self.text_bool(True))
        self.check('"abc" < "z"', self.text_bool(True))
        self.check('(3, 2) < (3, 4)', self.text_bool(True))
        self.check_not_optimized('%r < %r' % (b('abc'), u('abc')))
        self.check('2 > 2', self.text_bool(False))
        self.check('2 >= 2', self.text_bool(True))
        # may raise a BytesWarning exception
        self.check_not_optimized('%r == %r' % (b("abc"), u("abc")))

        # the type of x.__lt__(y) result can be different than bool,
        # not(x == y) cannot be replaced with x != y
        self.check_not_optimized("not(x == y)")
        self.check_not_optimized("not(x != y)")
        self.check_not_optimized("not(x <= y)")
        self.check_not_optimized("not(x >= y)")
        self.check_not_optimized("not(x < y)")
        self.check_not_optimized("not(x > y)")

    def test_Compare_is(self):
        self.check('True is True', self.text_bool(True))
        self.check('False is False', self.text_bool(True))
        self.check('None is None', self.text_bool(True))
        self.check('True is None', self.text_bool(False))
        self.check('True is not True', self.text_bool(False))
        self.check('True is not None', self.text_bool(True))

        # strings are not considered as singletons
        self.check_not_optimized('"" is not ""')
        self.check_not_optimized('"" is ""')

        # CPython implementation detail on small integers is ignored:
        # integers in range [-5; 256] are singletons
        self.check_not_optimized('1 is 1')
        self.check_not_optimized('1 is not 1')

        self.check("not(x in y)", self.text_ast("x not in y"))
        self.check("not(x not in y)", self.text_ast("x in y"))

    def test_Compare_in(self):
        config = self.create_config('builtin_funcs')

        # constant
        self.check('"xxx" in "abcdef"', self.text_bool(False))
        self.check('"xxx" not in "abcdef"', self.text_bool(True))
        self.check('"bcd" in "abcdef"', self.text_bool(True))
        self.check('"bcd" not in "abcdef"', self.text_bool(False))
        self.check('0 in (1, 2, 3)', self.text_bool(False))
        self.check('2 in (1, 2, 3)', self.text_bool(True))

        # generic
        self.check_not_optimized("x in (1, 2, 3)")
        self.check("x in [1, 2, 3]", self.text_ast('x in (1, 2, 3)'))
        self.check("x not in [1, 2, 3]", self.text_ast('x not in (1, 2, 3)'))

        if sys.version_info >= (3, 2):
            self.check("x in frozenset((1, 2, 3))", self.text_ast('x in {1, 2, 3}'), config)
        else:
            self.check_not_optimized("x in frozenset((1, 2, 3))", config)

        if PYTHON27:
            self.check("x in set((1, 2, 3))",
                       self.text_ast("x in {1, 2, 3}"),
                       config)
            self.check_not_optimized("x in {1, 2, 3}")
        else:
            self.check_not_optimized("x in set((1, 2, 3))", config)

        self.check("not(x in y)", self.text_ast("x not in y"))
        self.check("not(x not in y)", self.text_ast("x in y"))

        # may raise a BytesWarning exception
        self.check_not_optimized('%r in %r' % (b("abc"), u("abc")))
        self.check_not_optimized('%r not in %r' % (b("abc"), u("abc")))

    def test_Name(self):
        config = self.create_config()
        config.add_constant('debug', 0)
        self.check('debug', self.text_num(0), config)

    def test_print_function(self):
        if PYTHON3:
            self.check_not_optimized('print(1, 2)')
            self.check_not_optimized('print(1.5)')

            config = self.create_config('builtin_funcs')
            self.check('print(1, 2)', self.text_ast('print(%r)' % ('1 2',)), config)
            self.check('print(1.5)', self.text_ast('print(%r)' % ('1.5',)), config)
            self.check('print("x =", 3)', self.text_ast('print("x = 3")'), config)
            self.check_not_optimized('print(1, 2, sep="-")', config)
            self.check_not_optimized('print(1, 2, end="")', config)
            self.check_not_optimized('print(b"abc")', config)
        else:
            # optimization disabled by default
            self.check_not_optimized('from __future__ import print_function\nprint(1, 2)')
            self.check_not_optimized('from __future__ import print_function\nprint(1.5)')

            # optimized
            config = self.create_config('builtin_funcs')
            self.check('from __future__ import print_function\nprint(1, 2)',
                       self.text_ast('from __future__ import print_function\nprint(%r)' % ('1 2',)), config)
            self.check('from __future__ import print_function\nprint(1.5)', self.text_ast('from __future__ import print_function\nprint(%r)' % ('1.5',)), config)
            self.check('from __future__ import print_function\nprint("x =", 3)',
                       self.text_ast('from __future__ import print_function\nprint("x = 3")'), config)

            # not optimized
            self.check_not_optimized('from __future__ import print_function\nprint(u"unicode")', config)
            self.check_not_optimized('from __future__ import print_function\nprint(1, 2, sep="-")', config)
            self.check_not_optimized('from __future__ import print_function\nprint(1, 2, end="")', config)

    def test_print_statement(self):
        if PYTHON3:
            return self.skipTest('specific to Python 2')

        self.check_not_optimized('print')
        self.check('print 1, 2', self.text_ast('print %r' % ('1 2',)))
        self.check('print 1.5', self.text_ast('print %r' % ('1.5',)))
        self.check('print u"x =", 3', self.text_ast('print %r' % u('x = 3')))
        self.check_not_optimized('print u"", "\\xe9"')

        self.check('print 1, 2,', self.text_ast('print %r,' % ('1 2',)))
        self.check('print >>f, "x =", 2', self.text_ast('print >> f, %r' % ('x = 2',)))

    def test_call_builtin(self):
        config = self.create_config('builtin_funcs')

        self.check('set(tuple(list(iterable)))',
                   self.text_ast('set(iterable)'),
                   config)
        self.check('set(x for x in "abc")',
                   self.text_set('a', 'b', 'c'),
                   config)
        self.check_not_optimized('set((1, "b", (), []))',
                   config)

    def test_optimize_iter(self):
        config = self.create_config('builtin_funcs')

        # optimized
        self.check('[x for x in range(3)]',
                   self.text_ast('[0, 1, 2]'),
                   config)
        self.check('[x*2 for x in range(0, 2 ** 20, 2 ** 18)]',
                   self.text_ast('[0, 524288, 1048576, 1572864]'),
                   config)

        # not optimized
        self.check_not_optimized(
            'from re import *; [x for x in range(1000)]',
            config,
            warning=['Disable optimizations because of "from re import *"',
                     'Optimization missed because the builtin list is shadowed'])
        self.check_not_optimized(
            'list=4; [x for x in iterable]',
            config,
            warning='Optimization missed because the builtin list is shadowed')

        if PYTHON2:
            self.check(
                'from math import pow as list; [x for x in range(1000)]',
                self.text_ast('from math import pow as list; [x for x in xrange(1000)]'),
                config,
                warning='Optimization missed because the builtin list is shadowed')
            self.check('[x for x in range(1000)]',
                       self.text_ast('list(xrange(1000))'),
                       config)
            self.check(
                'from math import pow as xrange; [x for x in range(1000)]',
                self.text_ast('from math import pow as xrange; list(range(1000))'),
                config,
                warning='Optimization missed because the builtin xrange is shadowed')
            self.check_not_optimized(
                'from math import pow as list, floor as xrange; [x for x in range(1000)]',
                config,
                warning=['Optimization missed because the builtin xrange is shadowed',
                         'Optimization missed because the builtin list is shadowed'])
            self.check('[x*2 for x in range(1, 1000)]',
                       self.text_ast('[x*2 for x in xrange(1, 1000)]'),
                       config)
            self.check('[x*2 for x in range(0, 1000, 4)]',
                       self.text_ast('[x*2 for x in xrange(0, 1000, 4)]'),
                       config)
        else:
            # Python 3
            self.check_not_optimized(
                'from math import pow as list; [x for x in range(1000)]',
                config,
                warning='Optimization missed because the builtin list is shadowed')
            self.check('[x for x in range(1000)]',
                       self.text_ast('list(range(1000))'),
                       config)
            self.check_not_optimized(
                '[x*2 for x in range(1, 1000)]',
                config)
            self.check_not_optimized(
                '[x*2 for x in range(0, 1000, 4)]',
                config)

        # not optimized
        self.check('[x*2 for x in range(0, 2 ** 100)]',
                   self.text_ast('[x*2 for x in range(0, %s)]' % (2 ** 100,)),
                   config)
        self.check_not_optimized('[x*2 for x in range(1, 2, 3, 4)]',
                   config)
        self.check_not_optimized('[x*2 for x in range(1, 2, step=3)]',
                   config)
        self.check_not_optimized('[x*2 for x in range(0, 5, 0.1)]',
                   config)

    def test_For(self):
        # range => xrange
        self.check_not_optimized('for x in range(n): print(x)')

        config = self.create_config('builtin_funcs')
        config.unroll_limit = 0
        self.check('for x in range(3): print(x)',
                   self.text_ast('for x in (0, 1, 2): print(x)'),
                   config)
        self.check('for x in list("abc"): print(x)',
                   self.text_ast('for x in ("a", "b", "c"): print(x)'),
                   config)
        self.check_not_optimized('for x in tuple(iterable): print(x)', config)
        self.check_not_optimized('for x in list(iterable): print(x)', config)
        if PYTHON2:
            self.check('for x in range(1000): print(x)',
                       self.text_ast('for x in xrange(1000): print(x)'),
                       config)

        self.check('for x in [1, 2, 3]: print(x)',
                   self.text_ast('for x in (1, 2, 3): print(x)'),
                   config)
        self.check('for x in "": pass', self.TEXT_PASS)

        config = self.create_config()
        config.max_tuple_length = 2
        config.unroll_limit = 0
        self.check('for x in [1, 2]: print(x)',
                   self.text_ast('for x in (1, 2): print(x)'), config)
        self.check_not_optimized('for x in [1, 2, 3]: print(x)', config)

    def test_For_unroll(self):
        no_unroll = self.create_config()
        no_unroll.unroll_limit = 0
        self.check_not_optimized('for i in (1, 2, 3):\n print(i)',
                                 config=no_unroll)
        self.check('for i in (1, 2, 3):\n print(i)',
                   self.text_ast('i=1; print(i); i=2; print(i); i=3; print(i)'))
        self.check('for i in (1, 2, 3):\n pass',
                   self.text_ast('i=3'),
                   info=re.compile('Remove dead code: For'))

    def test_max_tuple_length(self):
        config = self.create_config('builtin_funcs')
        config.max_tuple_length = 3

        self.check('len((1, 2, 3))', self.text_num(3), config)
        self.check_not_optimized('len((1,2,3,4))', config)

        self.check('len(frozenset("abc"))', self.text_num(3), config)
        self.check('len(frozenset("abcabc"))', self.text_num(3), config)
        self.check_not_optimized('len(frozenset("dcba"))', config)

        self.check('set("abc")', self.text_set("a", "b", "c"), config)
        self.check('set("abcabc")', self.text_set("a", "b", "c"), config)
        self.check_not_optimized('set("abcd")', config)

    def test_iter_empty_iterable(self):
        config = self.create_config('builtin_funcs')
        self.check('dict(iter(()))', self.text_ast('{}'), config)
        self.check('frozenset(iter(()))', self.text_ast('frozenset()'), config)
        self.check('list(iter(()))', self.text_ast('[]'), config)
        self.check('set(iter(()))', self.text_set(), config)
        self.check('tuple(iter(()))', self.text_tuple(), config)

    def test_drop_empty_iterable(self):
        config = self.create_config('builtin_funcs')
        empty_set = self.text_set()
        self.check('set(())', empty_set, config)
        self.check('set([])', empty_set, config)
        self.check('set(tuple())', empty_set, config)
        self.check('set(list())', empty_set, config)
        self.check('set(dict())', empty_set, config)
        self.check('set(set())', empty_set, config)
        self.check('set(frozenset())', empty_set, config)

    def test_drop_builtin_iterable(self):
        config = self.create_config('builtin_funcs')
        self.check('dict(tuple(iterable))', self.text_ast('dict(iterable)'), config)
        self.check('frozenset(list(iterable))', self.text_ast('frozenset(iterable)'), config)
        self.check('set(tuple(iterable))', self.text_ast('set(iterable)'), config)
        self.check('tuple(list(iterable))', self.text_ast('tuple(iterable)'), config)
        self.check('tuple(tuple(iterable))', self.text_ast('tuple(iterable)'), config)
        self.check_not_optimized('tuple(frozenset(iterable))', config)
        self.check_not_optimized('tuple(set(iterable))', config)
        self.check_not_optimized('tuple(dict(iterable))', config)

    def test_drop_iter_empty_iterable(self):
        config = self.create_config('builtin_funcs')
        empty_tuple = self.text_tuple()
        self.check('tuple(iter(""))', empty_tuple, config)
        self.check('tuple(iter(()))', empty_tuple, config)
        self.check('tuple(iter([]))', empty_tuple, config)
        self.check('tuple(iter(tuple()))', empty_tuple, config)
        self.check('tuple(iter(list()))', empty_tuple, config)
        self.check('tuple(iter(dict()))', empty_tuple, config)
        self.check('tuple(iter(set()))', empty_tuple, config)
        self.check('tuple(iter(frozenset()))', empty_tuple, config)

    def test_GeneratorExp(self):
        config = self.create_config('builtin_funcs')

        # generators
        self.check('(x*2 for x in "")',
                   self.text_ast('(None for x in ())'))
        if PYTHON2:
            self.check('(x*2 for x in range(1000))',
                       self.text_ast('(x*2 for x in xrange(1000))'),
                       config)
        else:
            self.check_not_optimized('(x*2 for x in range(1000))', config)
        self.check('(x*2 for x in range(3))',
                   self.text_ast('(x*2 for x in (0, 1, 2))'),
                   config)

        # if
        self.check('(x*2 for x in (1, 2, 3) if True)',
                   self.text_ast('(x*2 for x in (1, 2, 3))'))
        self.check('(x*2 for x in (1, 2, 3) if 0)',
                   self.text_ast('(None for x in ())'))
        self.check('(x*2 for x in (1, 2, 3) if 0)',
                   self.text_ast('(None for x in ())'))
        self.check('(x*2 for x in (1, 2, 3) if True if 0 if 1)',
                   self.text_ast('(None for x in ())'))
        self.check_not_optimized('(x*2 for x in (1, 2, 3) if x % 2)')

        # generators
        self.check('(x * 2 for x in "")',
                   self.text_ast('(None for x in ())'))
        self.check('(x for x in "abc")',
                   self.text_ast('(x for x in ("a", "b", "c"))'))
        self.check_not_optimized('(x for x in iterable)')

        # tuple(generator)
        self.check('tuple(x for x in "abc")',
                   self.text_tuple("a", "b", "c"),
                   config)
        self.check('tuple(x*2 for x in "abc" if 0)',
                   self.text_tuple(),
                   config)
        self.check('tuple(x for x in iterable)',
                   self.text_ast('tuple(iterable)'),
                   config)

        # list(generator)
        self.check('list(x for x in range(3))',
                   self.text_ast('[0, 1, 2]'),
                   config)

    def test_ListComp(self):
        config = self.create_config('builtin_funcs')

        # list comprehension
        self.check('[x for x in "abc"]',
                   self.text_ast('["a", "b", "c"]'))
        self.check('[x for x in (0, 1, 2)]',
                   self.text_ast('[0, 1, 2]'))
        self.check('[x*2 for x in ""]',
                   self.text_ast('[]'))
        if PYTHON2:
            self.check('[x*2 for x in range(1000)]',
                       self.text_ast('[x*2 for x in xrange(1000)]'),
                       config)
        else:
            self.check_not_optimized('[x*2 for x in range(1000)]', config)
        self.check('[x*2 for x in range(3)]',
                   self.text_ast('[0, 2, 4]'),
                   config)

        # if
        self.check('[x*3 for x in (1, 2, 3) if True]',
                   self.text_ast('[3, 6, 9]'))
        self.check('[x*5 for x in (1, 2, 3) if True if 1]',
                   self.text_ast('[5, 10, 15]'))
        self.check('[x*2 for x in (1, 2, 3) if False]',
                   self.text_ast('[]'))
        self.check('[x*2 for x in (1, 2, 3) if 0 if True]',
                   self.text_ast('[]'))
        self.check_not_optimized('[x*2 for x in (1, 2, 3) if x % 2]')

        # replace with list()
        self.check('[x * 2 for x in ""]',
                   self.text_ast('[]'))
        self.check('[x for x in "abc"]',
                   self.text_ast('["a", "b", "c"]'))
        self.check('[x for x in "abc" if True]',
                   self.text_ast('["a", "b", "c"]'))
        self.check('[x for x in "abc" if False]',
                   self.text_ast('[]'))
        self.check_not_optimized('[x for x in iterable]')
        self.check('[x for x in iterable]',
                   self.text_ast('list(iterable)'),
                   config)

        # tuple(list comprehesion)
        self.check('tuple([x for x in "abc"])',
                   self.text_tuple("a", "b", "c"),
                   config)
        self.check('tuple([x for x in "abc" if False])',
                   self.text_tuple(),
                   config)
        self.check('tuple([x for x in iterable])',
                   self.text_ast('tuple(iterable)'),
                   config)
        if PYTHON2:
            self.check('set([x for x in range(1000)])',
                       self.text_ast('set(xrange(1000))'),
                       config)
        else:
            self.check('set([x for x in range(1000)])',
                       self.text_ast('set(range(1000))'),
                       config)

        self.check('set([x for x in range(3)])',
                   self.text_set(0, 1, 2),
                   config)

    def test_SetComp(self):
        if sys.version_info < (2, 7):
            return self.skipTest("need python 2.7+")
        config = self.create_config('builtin_funcs')

        # list comprehension
        self.check('{x for x in "abc"}',
                   self.text_set('a', 'b', 'c'))
        self.check('{x for x in (0, 1, 2)}',
                   self.text_set(0, 1, 2))
        self.check('{x*2 for x in ""}',
                   self.text_set())
        self.check('{x*2 for x in ""}',
                   self.text_set(),
                   config)
        if PYTHON2:
            self.check('{x*2 for x in range(1000)}',
                       self.text_ast('{x*2 for x in xrange(1000)}'),
                       config)
        else:
            self.check_not_optimized('{x*2 for x in range(1000)}', config)
        self.check('{x*2 for x in range(3)}',
                   self.text_ast('{x*2 for x in (0, 1, 2)}'),
                   config)

        # if
        self.check('{x*2 for x in (1, 2, 3) if True}',
                   self.text_ast('{x*2 for x in (1, 2, 3)}'))
        self.check('{x*2 for x in (1, 2, 3) if True if 1}',
                   self.text_ast('{x*2 for x in (1, 2, 3)}'))
        self.check('{x*2 for x in (1, 2, 3) if False}',
                   self.text_set())
        self.check('{x*2 for x in (1, 2, 3) if 0 if True}',
                   self.text_set())
        self.check_not_optimized('{x*2 for x in (1, 2, 3) if x % 2}')

        # replace with list()
        self.check('{x * 2 for x in ""}',
                   self.text_set())
        self.check('{x for x in "abc"}',
                   self.text_set('a', 'b', 'c'))
        self.check('{x for x in "abc" if True}',
                   self.text_set('a', 'b', 'c'))
        self.check('{x for x in "abc" if False}',
                   self.text_set())
        self.check_not_optimized('{x for x in iterable}')
        self.check('{x for x in iterable}',
                   self.text_ast('set(iterable)'),
                   config)

        # tuple(list comprehesion)
        self.check('tuple({x for x in "abc"})',
                   self.text_tuple("a", "b", "c"),
                   config)
        self.check('tuple({x for x in "abc" if False})',
                   self.text_tuple(),
                   config)
        self.check('tuple({x for x in iterable})',
                   self.text_ast('tuple(set(iterable))'),
                   config)
        if PYTHON2:
            self.check('set({x for x in range(1000)})',
                       self.text_ast('set(xrange(1000))'),
                       config)
        else:
            self.check('set({x for x in range(1000)})',
                       self.text_ast('set(range(1000))'),
                       config)

        self.check('set({x for x in range(3)})',
                   self.text_set(0, 1, 2),
                   config)

    def test_DictComp(self):
        if sys.version_info < (2, 7):
            return self.skipTest("need python 2.7+")
        config = self.create_config('builtin_funcs')

        # dict comprehension
        self.check('{key: value for key, value in {1: 2, 3: 4}}',
                   self.text_ast('{1: 2, 3: 4}'))
        self.check('{key: value for key, value in iterable}',
                   self.text_ast('dict(iterable)'),
                   config)
        self.check_not_optimized('{value: key for key, value in iterable}')

        # if
        self.check('{x: x*2 for x in (1, 2, 3) if True}',
                   self.text_ast('{x: x*2 for x in (1, 2, 3)}'))
        self.check('{x: x*2 for x in (1, 2, 3) if True if 1}',
                   self.text_ast('{x: x*2 for x in (1, 2, 3)}'))
        self.check('{x: x*2 for x in (1, 2, 3) if False}',
                   self.text_ast('{}'))
        self.check('{x: x*2 for x in (1, 2, 3) if 0 if True}',
                   self.text_ast('{}'))
        self.check_not_optimized('{x*2 for x in (1, 2, 3) if x % 2}')



class TestFrozenset(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.enable('builtin_funcs')
        return config

    def text_frozenset_str(self, text):
        arg = self._text_native_str(text)
        func = self._text_name('frozenset')
        call = 'Call(func=%s, args=[%s], keywords=[], starargs=None, kwargs=None)' % (func, arg)
        return self._text_module(call)

    def test_duplicates(self):
        self.check('frozenset("cba")', self.text_ast('frozenset(("c", "b", "a"))'))
        self.check('frozenset("abab")', self.text_ast('frozenset(("a", "b"))'))

    def test_BinOp(self):
        config = self.create_config()
        self.check_not_optimized('frozenset("ab") | frozenset("bc")', config)

        self.check('frozenset("ab") | frozenset("bc")',
                   self.text_frozenset_str('abc'))
        self.check('frozenset("ab") ^ frozenset("bc")',
                   self.text_frozenset_str("ac"))
        self.check('frozenset("abc") - frozenset("c")',
                   self.text_frozenset_str("ab"))
        self.check('frozenset("abc") & frozenset("bc")',
                   self.text_frozenset_str("bc"))

    def test_Compare(self):
        config = self.create_config()
        self.check_not_optimized('frozenset("abc") <= frozenset("abc")', config)
        self.check_not_optimized('"b" in frozenset("abc")', config)

        self.check('frozenset("abc") <= frozenset("abc")', self.text_bool(True))
        self.check('frozenset("abc") < frozenset("abc")', self.text_bool(False))
        self.check('"x" in frozenset("abc")', self.text_bool(False))
        self.check('"x" not in frozenset("abc")', self.text_bool(True))
        self.check('"b" in frozenset("abc")', self.text_bool(True))
        # raise a TypeError
        self.check_not_optimized('(1,) < frozenset()')
        self.check('(1, 2, 3) < frozenset("cba")',
                   self.text_ast('(1, 2, 3) < frozenset(("c", "b", "a"))'))

    def test_methods(self):
        config = self.create_config()
        self.check_not_optimized('frozenset("ab").union("c")', config)

        self.check('frozenset("ab").union("c")',
                   self.text_frozenset_str("abc"))
        self.check('frozenset("ax").union(frozenset("b"))',
                   self.text_frozenset_str("abx"))
        self.check('frozenset("ax").union(("b",))',
                   self.text_frozenset_str("abx"))
        self.check('frozenset("ab").intersection("bc")',
                   self.text_frozenset_str("b"))
        self.check('frozenset("ab").difference("bc")',
                   self.text_frozenset_str("a"))
        self.check('frozenset("ab").symmetric_difference("bc")',
                   self.text_frozenset_str("ac"))

        self.check('frozenset("abc").issubset(frozenset("bc"))',
                   self.text_bool(False))
        self.check('frozenset("abc").issuperset(frozenset("bc"))',
                   self.text_bool(True))
        self.check('frozenset("abc").isdisjoint(frozenset("bc"))',
                   self.text_bool(False))
        self.check('frozenset("abc").isdisjoint(frozenset("def"))',
                   self.text_bool(True))


# FIXME: remove these tests when assignment feature will be fixed
class TestAssignDisabled(BaseTestCase):
    def test_shadowed_builtin(self):
        config = self.create_config('builtin_funcs')
        self.check_not_optimized('len = ord; print(len("A"))', config)
        self.check_not_optimized('len += 1; print(len("A"))', config)
        self.check_not_optimized('del len; print(len("A"))', config)

    def test_shadowed_constants(self):
        config = self.create_config('builtin_funcs', 'pythonbin')
        self.check('print(sys.version)',
                   self.text_ast('print(%r)' % sys.version),
                   config)
        self.check_not_optimized('sys.version = "abc"; print(sys.version)',
                                 config)

class TestAssign(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.use_experimental_vars = True
        return config

    def test_shadowed_builtin(self):
        config = self.create_config('builtin_funcs')
        self.check_not_optimized('len = ord; print(len("A"))', config)
        self.check_not_optimized('len += 1; print(len("A"))', config)
        self.check_not_optimized('del len; print(len("A"))', config)

    def test_shadowed_constants(self):
        config = self.create_default_config()
        config.enable('builtin_funcs')
        config.enable('pythonbin')
        self.check('print(sys.version)',
                   self.text_ast('print(%r)' % sys.version),
                   config)
        self.check_not_optimized('sys.version = "abc"; print(sys.version)',
                                 config)

    def check_values(self, code, expected, unassigned=None, disabled=False, warning=None):
        tree = parse_ast(code)
        tree, optimizer = self._optimize_ast(tree)
        self.assertEqual(optimizer.namespace.values, expected)
        if unassigned is not None:
            self.assertEqual(optimizer.namespace._unassigned, set(unassigned))
        else:
            self.assertEqual(optimizer.namespace._unassigned, set())
        self.assertEqual(optimizer.namespace._vars_enabled, not disabled)
        self._check_warnings(warning)

    def test_supported(self):
        self.check_values('x=1',
                          {'x': 1})
        self.check_values('x,y=1,2',
                          {'x': 1, 'y': 2})
        self.check_values('a=b=c=3',
                          {'a': 3, 'b': 3, 'c': 3})
        self.check_values('a = b, c = (1, 2)',
                          {'a': (1, 2), 'b': 1, 'c': 2})

        self.check_values('x=1; x=2', {'x': 2})
        self.check_values('x,y=1,2; x,y=y,x', {'x': 2, 'y': 1})
        # FIXME: implement this
        #self.check_values('(x, y), z = ((1, 2), 3)',
        #                  {'x': 1, 'y': 2, 'z': 3})

        self.check('def f():\n x=math.e\n return x',
                   self.text_ast('def f():\n x=%r\n return %r'
                                 % (math.e, math.e)))
        self.check_not_optimized('''
def f():
    global x
    x = 1
    return x
        '''.strip())
        if PYTHON3:
            self.check_not_optimized('''
def f():
    x = 0
    def g():
        nonlocal x
        x = 1
        return x
    return g
            '''.strip())

    def test_local_namespace(self):
        # FunctionDef and ClassDef have a local namespace
        self.check('def f():\n x=1\n return x\nprint(x)',
                   self.text_ast('def f():\n x=1\n return 1\nprint(x)'))
        self.check('class Name:\n x=1\n y=x\nprint(x)',
                   self.text_ast('class Name:\n x=1\n y=1\nprint(x)'))

        code = '\n'.join((
            'class Name:',
            '    x = "abc"',
            '    y = x',
            '    class Nested:',
            '        print(y)',
            'print(x)',
        ))
        expected = code.replace('y = x', 'y = "abc"')
        self.check(code, self.text_ast(expected))

        code = '\n'.join((
            'class Name:',
            '    x = "abc"',
            '    def method(self):',
            '        x = 123',
            '        return x',
            '    print(x)',
        ))
        expected = '\n'.join((
            'class Name:',
            '    x = "abc"',
            '    def method(self):',
            '        x = 123',
            '        return 123',
            '    print("abc")',
        ))
        self.check(code, self.text_ast(expected))

    def test_unassigned(self):
        # unassign when assignment is not supported
        self.check_values('obj.attr = 1',
                          {}, ['obj.attr'])
        self.check_values('if test: trace = True',
                          {}, ['trace'])
        self.check_values('while test: trace = True',
                          {}, ['trace'])
        self.check_values('for x in range(test): trace = True',
                          {}, ['x', 'trace'])
        self.check_values('x=[]', {}, ['x'])
        self.check_values('x=[]; x=2', {}, ['x'])
        self.check_values('x=(1, 2); x[:1] = (2,)', {}, ('x',))
        self.check_values('x=1; x += 1', {}, ['x'])
        self.check_values('x=1; del x', {}, ['x'])
        self.check_values('x, y = x_y', {}, ['x', 'y'])
        self.check_values('x, y = func()', {}, ['x', 'y'])
        self.check_values('x = 1; a = x, y = c = 1', {'a': 1, 'c': 1}, ['x', 'y'])

    def test_while(self):
        self.check_not_optimized('i=0\nwhile i < 10: i = i + 1')

    def test_DictComp(self):
        if sys.version_info < (2, 7):
            return self.skipTest("need python 2.7+")
        self.check_not_optimized('k=1\nv=2\n{v:k for k, v in iterable}')

    def test_FunctionDef(self):
        self.check('x=1\ndef f():\n x=2\n return x',
                   self.text_ast('x=1\ndef f():\n x=2\n return 2'))
        self.check('def f(x=1+1): return x',
                   self.text_ast('def f(x=2): return x'))
        self.check_not_optimized('x=1\ndef f():\n return x')
        self.check_not_optimized('''
def f():
    a=1
    def g():
        print(a)
        '''.strip())

    def test_with(self):
        self.check_values('with open_file as x: pass',
                          {}, ['x'])
        self.check_values('x=1\nwith open_file as x: pass',
                          {}, ['x'])
        self.check_not_optimized('f=None\nwith open_file as f: f.write("abc")')

    def test_for(self):
        self.check_values('for x in range(3): print(x)',
                          {}, ['x'])
        self.check_values('for x, y in data: print(x)',
                          {}, ['x', 'y'])
        self.check_values('x=1\nfor x in range(3): print(x)',
                          {}, ['x'])
        self.check('def f():\n x = 5\n for x in "": x = 9\n return x',
                   self.text_ast('def f():\n x = 5\n return 5'))

    def test_listcomp(self):
        self.check_values('[x for x in range(3)]',
                          {}, ['x'])
        self.check_values('x=1\n[x for x in range(3)]',
                          {}, ['x'])

    def test_functiondef(self):
        self.check_not_optimized('x=1\ndef f(x): return x')

    def test_try(self):
        # FIXME: x=2
        self.check_values('try: x=1\nfinally: x=2',
                          {}, ['x'])
        self.check_values('try: x=1\nexcept: x=2',
                          {}, ['x'])
        self.check_values('try: x=1\nexcept: x=2\nelse: x=3',
                          {}, ['x'])
        self.check_values('try: import multiprocessing\nexcept ImportError: multiprocessing=None',
                          {}, ['multiprocessing'])

    def test_disabled(self):
        # unsupported: disable completly variables
        self.check_values('(x, y), z = x_y_z', {}, disabled=True,
                          warning=re.compile('^Disable optimizations on variables: '))

    def test_constant(self):
        config = self.create_config()
        config.add_constant('x', 42)
        self.check('x', self.text_num(42), config)
        self.check_not_optimized('x=1', config)

    def test_code(self):
        self.check('def f():\n x=1\n return x',
                   self.text_ast('def f():\n x=1\n return 1'))
        self.check('x="abcdef"; print(x[:3])',
                   self.text_ast('x="abcdef"; print("abc")'))


class TestBuiltinFuncs(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.enable('builtin_funcs')
        return config

    def test_optimize_iter_builtin(self):
        config = self.create_config()
        self.check_not_optimized('set([1, 2, 3])', config)

        self.check('frozenset([1, 2, 3])',
                   self.text_ast('frozenset((1, 2, 3))'))
        self.check('list([1, 2, 3])',
                   self.text_ast('[1, 2, 3]'))
        if PYTHON27:
            self.check('set([1, 2, 3])',
                       self.text_ast('{1, 2, 3}'))
        else:
            self.check('set([1, 2, 3])',
                       self.text_ast('set((1, 2, 3))'))

        self.check('dict({1: 2, 3: 4})',
                   self.text_ast('{1: 2, 3: 4}'))

    def test_no_builtin_funcs(self):
        config = self.create_config()
        self.check_not_optimized('abs(-5)', config)
        self.check_not_optimized('len("abc")', config)
        self.check_not_optimized('"a" in frozenset("abc")', config)

    def test_funcs(self):
        self.check('len("abc")', self.text_num(3))
        self.check('len("abcdef"[:2])', self.text_num(2))
        self.check('len((1,2,3))', self.text_num(3))
        self.check('len((1,2,3,4,5)[:2])', self.text_num(2))
        self.check('abs(-5)', self.text_num(5))
        self.check_not_optimized('abs("abc")')
        self.check('bin(5)', self.text_native_str('0b101'))
        self.check_not_optimized('bin(3.5)')
        self.check('bool(0)', self.text_bool(False))
        self.check('bool(3)', self.text_bool(True))
        self.check('bool("")', self.text_bool(False))
        self.check('bool("xxx")', self.text_bool(True))
        self.check('chr(65)', self.text_native_str('A'))
        self.check('divmod(10, 3)', self.text_ast('(3, 1)'))
        self.check('divmod(10.5, 3.0)', self.text_ast('(3.0, 1.5)'))
        self.check('float(6)', self.text_num(6.0))
        self.check('float(6.9)', self.text_num(6.9))
        self.check_not_optimized('float(1j)')
        self.check('int(6)', self.text_num(6))
        self.check('int(6.9)', self.text_num(6))
        self.check_not_optimized('int(1j)')
        self.check_not_optimized('oct(3.5)')

        self.check('ord(%r)' % b("A"), self.text_num(65))
        self.check('ord(%r)' % u("A"), self.text_num(65))
        self.check_not_optimized('ord("abc")')
        if sys.version_info >= (3,3):
            self.check('ord(%r)' % u("\U0010ffff"), self.text_num(0x10ffff))
        else:
            self.check_not_optimized('ord(%r)' % u("\U0010ffff"))
        config = self.create_config('builtin_funcs', 'pythonbin')
        if sys.maxunicode == 0x10ffff:
            self.check('ord(%r)' % u("\U0010ffff"), self.text_num(0x10ffff), config)
        else:
            self.check_not_optimized('ord(%r)' % u("\U0010ffff"), config)

        self.check("min('abc')", self.text_native_str('a'))
        self.check("min((1, 2, 3))", self.text_num(1))
        self.check("min(frozenset((1, 2, 3)))", self.text_num(1))
        self.check_not_optimized("min(1)")
        self.check("max('abc')", self.text_native_str('c'))
        self.check("max((1, 2, 3))", self.text_num(3))
        self.check("max(frozenset((1, 2, 3)))", self.text_num(3))
        self.check_not_optimized("max(1)")

        self.check('pow(2, 3)', self.text_num(8))
        self.check('pow(2, 123, 4567)', self.text_num(197))
        self.check('pow(10.0, 3)', self.text_num(1000.0))
        self.check_not_optimized('pow(-25.0, 0.5)')
        self.check_not_optimized('pow(10.0, 3, 7)')
        self.check('repr(1)', self.text_native_str('1'))
        self.check("repr('abc')", self.text_native_str("'abc'"))
        if PYTHON2:
            self.check('repr(1L)', self.text_native_str('1L'))
        self.check('str(1)', self.text_native_str('1'))
        self.check("str('abc')", self.text_native_str('abc'))
        if PYTHON3:
            self.check_not_optimized("str(%r)" % b('abc'))
        else:
            self.check("str(u'abc')", self.text_bytes('abc'))
            self.check_not_optimized("str(u'\xe9')")
        self.check("sum((1, 2, 3))", self.text_num(6))
        self.check("sum(frozenset((1, 2, 3)))", self.text_num(6))
        self.check("sum((1, 2, 3), 0)", self.text_num(6))
        self.check_not_optimized("sum(((1,), 2))")
        self.check_not_optimized("sum((1, 2), 'a')")
        self.check("tuple()", self.text_tuple())
        self.check("tuple('abc')", self.text_tuple('a', 'b', 'c'))
        self.check("tuple((1, 2, 3.0))", self.text_tuple(1, 2, 3.0))

        if PYTHON3:
            self.check('oct(5)', self.text_unicode('0o5'))
            self.check("round(5)", self.text_num(5))
            self.check("round(5.5)", self.text_num(6))
            self.check("round(5, 1)", self.text_num(5))
            self.check("round(5.567, 1)", self.text_num(5.6))
        else:
            self.check('oct(5)', self.text_native_str('05'))
            self.check('long(6)', self.text_num(long(6)))
            self.check('long(6.9)', self.text_num(long(6)))
            self.check_not_optimized('long(1j)')
            self.check('unichr(65)', self.text_unicode(u('A')))
            self.check('repr(1L)', self.text_native_str('1L'))
            self.check('str(1L)', self.text_native_str('1'))
            self.check('unicode(1)', self.text_unicode(u('1')))
            self.check('unicode(1L)', self.text_unicode(u('1')))
            self.check("unicode('abc')", self.text_unicode(u('abc')))
            self.check_not_optimized("unicode('abc\xff')")
            self.check("round(5)", self.text_num(5.0))
            self.check("round(5.5)", self.text_num(6.0))
            self.check("round(5, 1)", self.text_num(5.0))
            self.check("round(5.567, 1)", self.text_num(5.6))

        config = self.create_config('builtin_funcs')
        config.min_int = -10
        config.max_int = 10
        config.max_string_length = 5
        self.check('pow(2, 3)', self.text_num(8), config)
        self.check_not_optimized('pow(2, 10)', config)
        self.check('pow(2, 10, 7)', self.text_num(2), config)
        self.check_not_optimized('pow(2, 10, 42)', config)
        self.check('repr(12345)', self.text_native_str('12345'), config)
        self.check_not_optimized('repr(1234567)', config)
        self.check('str(12345)', self.text_native_str('12345'), config)
        self.check_not_optimized('str(1234567)', config)

        config = self.create_config('builtin_funcs')
        config.max_size = 3
        self.check('len((1,2,3))', self.text_num(3), config)
        self.check_not_optimized('len((1,2,3,4))', config)

        self.check('float("6")', self.text_num(6.0))
        self.check('float("6.9")', self.text_num(6.9))
        self.check('float("inf")', self.text_num(float("inf")))
        self.check('float("-inf")', self.text_num(float("-inf")))
        self.check('float("nan")', self.text_num(float("nan")))
        self.check('float("-0.0")', self.text_num(float("-0.0")))
        self.check_not_optimized('float("abc")')

        config = self.create_config('builtin_funcs')
        config.max_tuple_length = 3
        self.check('tuple("abc")', self.text_ast('("a", "b", "c")'), config)
        self.check_not_optimized('tuple("abcd")', config)

        if sys.version_info >= (3,3):
            self.check('chr(0xffff)', self.text_unicode(chr(0xffff)))
            self.check('chr(0x10ffff)', self.text_unicode(chr(0x10ffff)))
            self.check_not_optimized('chr(0x110000)')
        elif PYTHON3:
            self.check('chr(0xffff)', self.text_unicode(chr(0xffff)))
            self.check_not_optimized('chr(0x10ffff)')
            config = self.create_config('builtin_funcs', 'pythonbin')
            self.check('chr(0x10ffff)', self.text_unicode(chr(0x10ffff)), config)
            self.check_not_optimized('chr(0x110000)')
        else:
            self.check_not_optimized('chr(0x100)')
            self.check('unichr(0xffff)', self.text_unicode(unichr(0xffff)))
            self.check_not_optimized('unichr(0x10ffff)')
            if sys.maxunicode > 0xffff:
                config = self.create_config('builtin_funcs', 'pythonbin')
                self.check('unichr(0x10ffff)', self.text_unicode(unichr(0x10ffff)), config)
            self.check_not_optimized('unichr(0x110000)')

class TestBuiltinTypes(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.enable('builtin_types')
        return config

    def test_bytes_methods(self):
        self.check('%r.decode("ascii")' % b("abc"), self.text_unicode(u('abc')))
        self.check('%r.decode("ascii", "replace")' % b("ab\xff"), self.text_unicode(u('ab\ufffd')))
        if sys.version_info < (3,3):
            self.check_not_optimized('%r.decode("utf-8")' % b('\xf4\x8f\xbf\xbf'))
        else:
            self.check(
                '%r.decode("utf-8")' % b('\xf4\x8f\xbf\xbf'),
                self.text_unicode(u('\U0010ffff')))

        self.check('%r.count(%r)' % (b('abab'), b('a')), self.text_num(2))
        self.check('%r.find(%r)' % (b('abc'), b('b')), self.text_num(1))
        self.check('%r.index(%r)' % (b('abc'), b('b')), self.text_num(1))
        self.check('%r.startswith(%r)' % (b('abc'), b('ab')), self.text_bool(True))
        self.check('%r.endswith(%r)' % (b('abc'), b('xx')), self.text_bool(False))

        self.check('%r.zfill(5)' % b('abc'), self.text_bytes(b('00abc')))
        self.check('%r.ljust(5)' % b('abc'), self.text_bytes(b('abc  ')))
        self.check('%r.rjust(5)' % b('abc'), self.text_bytes(b('  abc')))
        self.check('%r.center(5, %r)' % (b('abc'), b('_')), self.text_bytes(b('_abc_')))
        self.check_not_optimized('%r.center(5, %r)' % (b('abc'), b('__')))

        self.check('%r.strip()' % b(' [x] '), self.text_bytes(b('[x]')))
        self.check('%r.lstrip()' % b(' [x] '), self.text_bytes(b('[x] ')))
        self.check('%r.rstrip()' % b(' [x] '), self.text_bytes(b(' [x]')))
        self.check('%r.strip(%r)' % (b('x[x]x'), b('x')), self.text_bytes(b('[x]')))
        self.check('%r.lstrip(%r)' % (b('x[x]x'), b('x')), self.text_bytes(b('[x]x')))
        self.check('%r.rstrip(%r)' % (b('x[x]x'), b('x')), self.text_bytes(b('x[x]')))

        self.check('%r.replace(%r, %r)' % (b('Hello World!'), b(' '), b('_')), self.text_bytes(b('Hello_World!')))

        self.check('%r.partition(%r)' % (b('user@host@domain'), b('@')),
                   self.text_tuple(b('user'), b('@'), b('host@domain')))
        self.check('%r.rpartition(%r)' % (b('user@host@domain'), b('@')),
                   self.text_tuple(b('user@host'), b('@'), b('domain')))

        if PYTHON2:
            self.check('%r.join(%r)' % (b('-'), b('abc')),
                       self.text_bytes(b('a-b-c')))
        else:
            self.check_not_optimized('%r.join(%r)' % (b('-'), b('abc')))
        self.check('%r.join(%r)' % (b(', '), (b('1'), b('5'), b('10'))),
                   self.text_bytes(b('1, 5, 10')))
        self.check_not_optimized('%r.join(%r)' % (b('-'), (b('a'), u('bc'))))
        self.check_not_optimized('%r.join(%r)' % (b('-'), (b('a'), 2)))

        self.check('%r.expandtabs()' % b("a\tb"), self.text_bytes(b("a       b")))
        self.check('%r.expandtabs(3)' % b("a\tb"), self.text_bytes(b("a  b")))

        if PYTHON3:
            self.check('%r.lower()' % b("abcDEF"), self.text_bytes(b("abcdef")))
            self.check('%r.upper()' % b("abcDEF"), self.text_bytes(b("ABCDEF")))
            self.check('%r.title()' % b("abcDEF"), self.text_bytes(b("Abcdef")))
            self.check('%r.swapcase()' % b("abcDEF"), self.text_bytes(b("ABCdef")))
            self.check('%r.isalnum()' % b("abc"), self.text_bool(True))
            self.check('%r.isalnum()' % b("-"), self.text_bool(False))
            self.check('%r.isalpha()' % b("abc"), self.text_bool(True))
            self.check('%r.isalpha()' % b("-"), self.text_bool(False))
            self.check('%r.isdigit()' % b("123"), self.text_bool(True))
            self.check('%r.isdigit()' % b("abc"), self.text_bool(False))
            self.check('%r.islower()' % b("abc"), self.text_bool(True))
            self.check('%r.islower()' % b("ABC"), self.text_bool(False))
            self.check('%r.isspace()' % b(" "), self.text_bool(True))
            self.check('%r.isspace()' % b("abc"), self.text_bool(False))
            self.check('%r.istitle()' % b("Abc"), self.text_bool(True))
            self.check('%r.istitle()' % b("ABC"), self.text_bool(False))
            self.check('%r.isupper()' % b("ABC"), self.text_bool(True))
            self.check('%r.isupper()' % b("Abc"), self.text_bool(False))

    def test_unicode_methods(self):
        self.check('%r.encode("ASCII")' % u('abc'), self.text_bytes(b('abc')))
        self.check('%r.encode("ASCII", "ignore")' % u('ab\xff'), self.text_bytes(b('ab')))
        self.check('%r.encode("ASCII", "backslashreplace")' % u('ab\xff'), self.text_bytes(b('ab\\xff')))
        self.check('%r.encode("latin1")' % u('abc'), self.text_bytes(b('abc')))
        if sys.version_info < (3,3):
            self.check_not_optimized('%r.encode("utf-8")' % u('\U0010ffff'))
        else:
            self.check(
                '%r.encode("utf-8")' % u('\U0010ffff'),
                self.text_bytes(b('\xf4\x8f\xbf\xbf')))

        self.check('%r.count(%r)' % (u('abab'), u('a')), self.text_num(2))
        self.check('%r.find(%r)' % (u('abcb'), u('b')), self.text_num(1))
        self.check('%r.rfind(%r)' % (u('abcb'), u('b')), self.text_num(3))
        self.check('%r.index(%r)' % (u('abcb'), u('b')), self.text_num(1))
        self.check('%r.rindex(%r)' % (u('abcb'), u('b')), self.text_num(3))
        self.check('%r.startswith(%r)' % (u('abc'), u('x')), self.text_bool(False))
        self.check('%r.endswith(%r)' % (u('abc'), u('bc')), self.text_bool(True))
        self.check_not_optimized('%r.index(%r)' % (u('abc'), u('xx')))

        self.check('%r.zfill(5)' % u('abc'), self.text_unicode(u('00abc')))
        self.check('%r.ljust(5)' % u('abc'), self.text_unicode(u('abc  ')))
        self.check('%r.rjust(5)' % u('abc'), self.text_unicode(u('  abc')))
        self.check('%r.center(5, %r)' % (u('abc'), u('_')), self.text_unicode(u('_abc_')))
        self.check_not_optimized('%r.center(5, %r)' % (u('abc'), u('__')))

        self.check('%r.strip()' % u(' [x] '), self.text_unicode(u('[x]')))
        self.check('%r.lstrip()' % u(' [x] '), self.text_unicode(u('[x] ')))
        self.check('%r.rstrip()' % u(' [x] '), self.text_unicode(u(' [x]')))
        self.check('%r.strip(%r)' % (u('x[x]x'), u('x')), self.text_unicode(u('[x]')))
        self.check('%r.lstrip(%r)' % (u('x[x]x'), u('x')), self.text_unicode(u('[x]x')))
        self.check('%r.rstrip(%r)' % (u('x[x]x'), u('x')), self.text_unicode(u('x[x]')))

        self.check('%r.replace(%r, %r)' % (u('Hello World!'), u(' '), u('_')), self.text_unicode(u('Hello_World!')))

        self.check('%r.partition(%r)' % (u('user@host@domain'), u('@')),
                   self.text_tuple(u('user'), u('@'), u('host@domain')))
        self.check('%r.rpartition(%r)' % (u('user@host@domain'), u('@')),
                   self.text_tuple(u('user@host'), u('@'), u('domain')))

        self.check('%r.join(%r)' % (u('-'), u('abc')),
                   self.text_unicode(u('a-b-c')))
        self.check('%r.join(%r)' % (u(', '), (u('1'), u('5'), u('10'))),
                   self.text_unicode(u('1, 5, 10')))
        self.check_not_optimized('%r.join(%r)' % (u('-'), (u('a'), b('bc'))))
        self.check_not_optimized('%r.join(%r)' % (u('-'), (u('a'), 2)))

        self.check('%r.expandtabs()' % u("a\tb"), self.text_unicode(u("a       b")))
        self.check('%r.expandtabs(3)' % u("a\tb"), self.text_unicode(u("a  b")))

        self.check('%r.lower()' % u("abcDEF"), self.text_unicode(u("abcdef")))
        self.check('%r.upper()' % u("abcDEF"), self.text_unicode(u("ABCDEF")))
        self.check('%r.title()' % u("abcDEF"), self.text_unicode(u("Abcdef")))
        self.check('%r.swapcase()' % u("abcDEF"), self.text_unicode(u("ABCdef")))
        self.check('%r.isalnum()' % u("abc"), self.text_bool(True))
        self.check('%r.isalnum()' % u("-"), self.text_bool(False))
        self.check('%r.isalpha()' % u("abc"), self.text_bool(True))
        self.check('%r.isalpha()' % u("-"), self.text_bool(False))
        self.check('%r.isdigit()' % u("123"), self.text_bool(True))
        self.check('%r.isdigit()' % u("abc"), self.text_bool(False))
        self.check('%r.islower()' % u("abc"), self.text_bool(True))
        self.check('%r.islower()' % u("ABC"), self.text_bool(False))
        self.check('%r.isspace()' % u(" "), self.text_bool(True))
        self.check('%r.isspace()' % u("abc"), self.text_bool(False))
        self.check('%r.istitle()' % u("Abc"), self.text_bool(True))
        self.check('%r.istitle()' % u("ABC"), self.text_bool(False))
        self.check('%r.isupper()' % u("ABC"), self.text_bool(True))
        self.check('%r.isupper()' % u("Abc"), self.text_bool(False))

    def test_tuple_methods(self):
        self.check('(2,0,2,2).count(2)', self.text_num(3))
        self.check('(2,0,2,2).index(0)', self.text_num(1))
        self.check_not_optimized('(2,0,2,2).index(9)')

    def test_int_methods(self):
        if PYTHON27:
            self.check('(12345).bit_length()', self.text_num(14))

    def test_float_methods(self):
        self.check('(1.0).is_integer()', self.text_bool(True))
        self.check('(1.3).is_integer()', self.text_bool(False))
        self.check('(1.5).as_integer_ratio()', self.text_tuple(3, 2))
        self.check('(1.5).hex()', self.text_native_str('0x1.8000000000000p+0'))

    def test_float_static_methods(self):
        self.check_not_optimized('float.fromhex("0x1.8p+0")')

        config = self.create_default_config()
        config.enable('builtin_funcs')
        self.check('float.fromhex("0x1.8p+0")', self.text_num(1.5), config)


class TestFeatures(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        # For float("inf") above
        return config

    def test_platform(self):
        import errno
        import os
        import platform
        import socket

        self.check_not_optimized('os.O_RDONLY')
        self.check_not_optimized('errno.ENOENT')
        self.check_not_optimized('sys.platform')

        config = self.create_config('platform')
        self.check('os.O_RDONLY', self.text_num(os.O_RDONLY), config)
        self.check('errno.ENOENT', self.text_num(errno.ENOENT), config)
        if PYTHON34:
            # Constants are now enum
            self.check('socket.AF_INET', self.text_num(socket.AF_INET.value), config)
            self.check('socket.SOCK_STREAM', self.text_num(socket.SOCK_STREAM.value), config)
        else:
            self.check('socket.AF_INET', self.text_num(socket.AF_INET), config)
            self.check('socket.SOCK_STREAM', self.text_num(socket.SOCK_STREAM), config)
        self.check('sys.platform', self.text_native_str(sys.platform), config)
        self.check('platform.machine()', self.text_native_str(platform.machine()), config)
        self.check('platform.release()', self.text_native_str(platform.release()), config)
        self.check('platform.system()', self.text_native_str(platform.system()), config)
        self.check('socket.ntohl(12)', self.text_num(socket.ntohl(12)), config)
        self.check('socket.ntohs(12)', self.text_num(socket.ntohs(12)), config)
        self.check('socket.htonl(12)', self.text_num(socket.htonl(12)), config)
        self.check('socket.htons(12)', self.text_num(socket.htons(12)), config)

    def test_pythonbin(self):
        import platform

        self.check_not_optimized('sys.hexversion')
        self.check_not_optimized('sys.platform')
        self.check_not_optimized('platform.python_version')
        if sys.version_info < (3,3):
            self.check_not_optimized('%r[-1]' % u("abc\U0010ffff-"))

        config = self.create_config('pythonbin')
        self.assertIn('platform', config.features)
        self.check('sys.hexversion', self.text_num(sys.hexversion), config)
        self.check('sys.version[:3]', self.text_native_str(sys.version[:3]), config)

        if PYTHON27:
            self.check_not_optimized('sys.version_info', config)
            self.check('sys.version_info.minor', self.text_num(sys.version_info.minor), config)
        else:
            self.check('sys.version_info', self.text_tuple(*sys.version_info), config)

        self.check('platform.python_version()', self.text_native_str(platform.python_version()), config)
        self.check('platform.python_version_tuple()', self.text_tuple(*platform.python_version_tuple()), config)
        if sys.version_info < (3,3):
            if sys.maxunicode == 0x10ffff:
                self.check(
                    '%r[3]' % u("abc\U0010ffff-"),
                    self.text_unicode(u("\U0010ffff")),
                    config)
            else:
                self.check(
                    '%r[3]' % u("abc\U0010ffff-"),
                    self.text_unicode(u("\udbff")),
                    config)
                self.check(
                    '%r[4]' % u("abc\U0010ffff-"),
                    self.text_unicode(u("\udfff")),
                    config)
            self.check(
                '%r[-1]' % u("abc\U0010ffff-"),
                self.text_unicode(u("-")),
                config)

    def test_pythonenv(self):
        self.check_not_optimized('__debug__')
        if PYTHON2:
            self.check_not_optimized('21 / 5')

        config = self.create_config('pythonenv')
        self.assertIn('platform', config.features)
        self.assertIn('pythonbin', config.features)
        self.check('__debug__', self.text_bool(__debug__), config)
        self.check_not_optimized('sys.flags', config)
        self.check('sys.flags.optimize', self.text_num(sys.flags.optimize), config)
        self.check('21 / 5', self.text_num(21 / 5), config)

    def test_math(self):
        self.check('math.e', self.text_num(math.e))
        self.check('math.pi', self.text_num(math.pi))
        self.check('math.exp(1.0)', self.text_num(math.e))
        self.check('math.log(1.0)', self.text_num(0.0))
        self.check('math.log(16.0, 2.0)', self.text_num(4.0))
        self.check('math.log(100.0, 10.0)', self.text_num(2.0))
        self.check('math.log(math.e)', self.text_num(1.0))
        if sys.version_info >= (3, 3):
            self.check('math.log2(32)', self.text_num(5.0))
        self.check('math.log10(1000)', self.text_num(3.0))
        self.check('math.pow(2, 4)', self.text_num(16.0))
        self.check('math.pow(10.0, 3)', self.text_num(1000.0))
        self.check_not_optimized('math.pow(2j, 4)')
        self.check_not_optimized('math.pow(-25.0, 0.5)')

        self.check('math.sqrt(4.0)', self.text_num(2.0))
        self.check_not_optimized('math.sqrt(-1)')
        if PYTHON3:
            self.check('math.ceil(1.5)', self.text_num(2))
            self.check('math.floor(1.5)', self.text_num(1))
        else:
            self.check('math.ceil(1.5)', self.text_num(2.0))
            self.check('math.floor(1.5)', self.text_num(1.0))
        self.check('math.copysign(1.0, -2.0)', self.text_num(-1.0))
        self.check('math.factorial(3)', self.text_num(6))
        self.check('math.fsum((1.0, 2.0, 4.0))', self.text_num(7.0))
        self.check('math.log1p(math.e - 1.0)', self.text_num(1.0))
        self.check('math.trunc(-1.5)', self.text_num(-1))
        self.check('math.trunc(4.5)', self.text_num(4))
        self.check('math.fabs(-2.0)', self.text_num(2.0))
        self.check('math.fmod(3.5, 2.0)', self.text_num(1.5))
        self.check('math.frexp(1.5)', self.text_tuple(0.75, 1))
        if sys.version_info >= (3, 2):
            config = self.create_default_config()
            config.enable('builtin_funcs')
            self.check('math.expm1(0.0)', self.text_num(0.0))
            self.check('math.isfinite(1.0)', self.text_bool(True))
            self.check('math.isfinite(float("inf"))', self.text_bool(False), config)
            self.check('math.isfinite(float("nan"))', self.text_bool(False), config)
            self.check('math.isinf(1.0)', self.text_bool(False))
            self.check('math.isinf(float("inf"))', self.text_bool(True), config)
            self.check('math.isinf(float("nan"))', self.text_bool(False), config)
            self.check('math.isnan(1.0)', self.text_bool(False))
            self.check('math.isnan(float("inf"))', self.text_bool(False), config)
            self.check('math.isnan(float("nan"))', self.text_bool(True), config)
        self.check('math.ldexp(0.75, 1)', self.text_num(1.5))
        self.check('math.modf(1.5)', self.text_tuple(0.5, 1.0))
        self.check('math.modf(-1.5)', self.text_tuple(-0.5, -1.0))

    def test_math_trigonometric(self):
        self.check('math.cos(0.0)', self.text_num(1.0))
        self.check('math.cos(math.pi)', self.text_num(-1.0))
        self.check('math.sin(0.0)', self.text_num(0.0))
        self.check('math.tan(0.0)', self.text_num(0.0))
        self.check('math.tan(0.0)', self.text_num(0.0))

        self.check('math.acos(1.0)', self.text_num(0.0))
        self.check('math.acos(-1.0)', self.text_num(math.pi))
        self.check('math.asin(0.0)', self.text_num(0.0))
        self.check('math.atan(0.0)', self.text_num(0.0))
        self.check('math.atan(0.0)', self.text_num(0.0))
        self.check('math.atan2(0.0, 1.0)', self.text_num(0.0))

        self.check('math.hypot(1.0, 1.0)', self.text_num(math.sqrt(2)))

    def test_string_module(self):
        self.check('string.ascii_lowercase',
                   self.text_native_str('abcdefghijklmnopqrstuvwxyz'))
        if PYTHON2:
            self.check('string.atof("1.0")', self.text_num(1.0))
            self.check_not_optimized('string.atof("xxx")')
            self.check('string.atoi("2")', self.text_num(2))
            self.check('string.atoi("10", 16)', self.text_num(16))
            self.check_not_optimized('string.atoi("xxx")')
            self.check('string.atol("2")', self.text_num(long(2)))
            self.check('string.atol("12", 8)', self.text_num(long(10)))
            self.check_not_optimized('string.atol("xxx")')
            self.check('string.expandtabs(%r)' % b("a\tb"), self.text_bytes(b("a       b")))
            self.check('string.expandtabs(%r)' % u("a\tb"), self.text_unicode(u("a       b")))
            self.check('string.expandtabs("a\tb", 3)', self.text_native_str("a  b"))
        if PYTHON3:
            self.check('string.capwords("abc def.")',
                       self.text_native_str('Abc Def.'))

    def test_Import(self):
        self.check_qualnames('import math as M', {'M': 'math'})
        self.check('import math as M; x=M.sqrt(4.0)',
                   self.text_ast('import math as M; x=2.0'))
        code = '\n'.join((
            'try:',
            '    import StringIO as BUFFER',
            'except ImportError:',
            '    import cStringIO as BUFFER',
        ))
        self.check_qualnames(code, {})

        self.check('from math import *\ndef f(): return pow(2, 4)',
                   self.text_ast('from math import *\ndef f(): return 16.0'))
        self.check_not_optimized('from unknown import *\ndef f(): return pow(2, 4)',
                                 warning='Disable optimizations because of "from unknown import *"')

        config = self.create_config('builtin_funcs')
        self.check('from math import *\ndef f(): return len("abc")',
                   self.text_ast('from math import *\ndef f(): return 3'),
                   config)
        self.check_not_optimized('from unknown import *\ndef f(): return len("abc")',
                                 config,
                                 warning='Disable optimizations because of "from unknown import *"')

    def test_ImportFrom(self):
        self.check_qualnames('from math import trunc as M', {'M': 'math.trunc'})
        self.check_qualnames('from math import pow; from math import pow',
                             {'pow': 'math.pow'})
        config = self.create_config('builtin_funcs')
        self.check('from re import *\nfrom math import pow\ndef f(): return pow(2, 4)',
                   self.text_ast('from re import *\nfrom math import pow\ndef f(): return 16.0'),
                   config,
                   warning='Disable optimizations because of "from re import *"')
        code = '\n'.join((
            'try:',
            '    from StringIO import StringIO as BUFFER',
            'except ImportError:',
            '    from cStringIO import StringIO as BUFFER',
        ))
        self.check_qualnames(code, {})

        # builtins are disabled
        config = self.create_config('builtin_funcs')
        self.check('from math import *; print(len("abc"))',
                   self.text_ast('from math import *; print("3")'),
                   config)
        self.check_not_optimized(
            'from re import *; print(len("abc"))',
            config,
            warning='Disable optimizations because of "from re import *"')


class TestREADME(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.enable('builtin_funcs')
        config.enable('pythonenv')
        config._constants['__debug__'] = True
        config._constants['sys.flags.optimize'] = 0
        config._constants['sys.maxunicode'] = 0x10ffff
        config._constants['sys.platform'] = 'linux2'
        config._constants['sys.maxint'] = 2147483647
        return config

    def check_line(self, line_number, line):
        if ('`' not in line) or ('=>' not in line):
            return
        line = line.replace(r'\\', '\\')
        parts = line.split('``')
        parts = parts[1:]
        before = parts[0]
        after = parts[2]
        after2 = None
        info = None

        if (sys.version_info < (2, 7)
        and ('.bit_length()' in before
             or 'sys.version_info.major' in before
             or after == '{1, 2, 3}')):
            return self.skipTest("need python 2.7+")
        elif (PYTHON3
        and 'string.atoi' in before):
            return self.skipTest("specific to Python 2")
        if PYTHON3 and before.startswith('u"'):
            before = before[1:]
        if before == 'if a: if b: print("true")':
            before = 'if a:\n if b:\n  print("true")'
        if after == 'if a and b: print("true")':
            after = 'if a and b:\n print("true")'
        if 'return 1; return 2' in before:
            info = re.compile('Remove unreachable code')
        elif before in ('for i in (1, 2, 3): pass',
                        'def f(a, b): s = a+b; 3; return s'):
            info = re.compile('Remove dead code')

        if 'if DEBUG:' in before:
            config = self.create_default_config()
            config.add_constant('DEBUG', False)
        else:
            config = None
        if len(parts) >= 5 and ' or ' in parts[3]:
            after2 = parts[4]
        try:
            self.check(before, self.text_ast(after), config, info=info)
        except (AssertionError, SyntaxError):
            if after2 is not None:
                self.check(before, self.text_ast(after2), config)
            else:
                raise

    def test_readme_file(self):
        filename = os.path.join(os.path.dirname(__file__), '..', 'README.rst')
        with open(filename) as fp:
            for line_number, line in enumerate(fp, 1):
                self.check_line(line_number, line)


class TestStruct(BaseTestCase):
    def create_default_config(self):
        config = BaseTestCase.create_default_config(self)
        config.enable('struct')
        return config

    def test_pack(self):
        self.check('struct.pack("B", 0x12)', self.text_bytes(b('\x12')))
        self.check_not_optimized('struct.pack("ii")')
        self.check_not_optimized('struct.pack("ii", 0x12)')

    def test_unpack(self):
        self.check('struct.unpack("B", %r)' % b('\x12'), self.text_tuple(0x12))
        self.check_not_optimized('struct.unpack("ii")')
        self.check_not_optimized('struct.unpack("ii", 1)')

    def test_calcsize(self):
        self.check('struct.calcsize("BBB")', self.text_num(3))
        self.check_not_optimized('struct.calcsize("#_")')


class TestNamespace(TestCase):
    def test_module_aliases(self):
        ns = Namespace()
        ns.add_alias('math', 'xxx')
        self.assertEqual(ns.get_qualname('xxx'), 'math')

        ns2 = Namespace()
        ns2.add_alias('math.trunc', 'trunc')
        self.assertEqual(ns2.get_qualname('trunc'), 'math.trunc')

        ns3 = Namespace()
        ns3.add_alias('os', 'O')
        ns3.add_alias('os.path', 'P')
        ns3.add_alias('os.path.exists', 'E')
        self.assertEqual(ns3.get_qualname('O'), 'os')
        self.assertEqual(ns3.get_qualname('P'), 'os.path')
        self.assertEqual(ns3.get_qualname('E'), 'os.path.exists')


def main():
    unittest.main(module=__name__)

if __name__ == "__main__":
    main()

