import pytest

from ._util import refs_equal, skip_type_union, skip_walrus


class TestFunction:
    @refs_equal
    def test_func_uses(self):
        s = 'import a\ndef f():\n  a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_async_func_uses(self):
        s = 'import a\nasync def f():\n  a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_name_shadows(self):
        s = 'import a\ndef a():\n  pass\n  a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_name_shadows_inside(self):
        s = 'import a\ndef a():\n  a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_assigns_then_uses(self):
        s = 'import a\ndef f():\n  a = 1\n  a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_assigns_then_used_outside(self):
        s = 'import a\ndef f():\n  a = 1\na'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_annotates_then_uses(self):
        s = 'import a\ndef f(arg: a):\n  arg.b'
        refs = [('a', 'a'), ('a', 'a'), ('a.()', 'arg'), ('a.().b', 'arg.b')]
        return s, refs

    @refs_equal
    def test_func_annotates_then_assigns(self):
        # Note: inner nodes after outer ones
        s = 'import a\ndef f(arg: a) -> a:\n  a = 1'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a'), ('a.()', 'arg')]
        return s, refs

    @refs_equal
    def test_func_annotates_as_generic_then_uses(self):
        s = 'import a\ndef f(arg: a[0]):\n  arg.b'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_annotates_inside_generic_then_uses(self):
        s = 'import a\ndef f(arg: b[a]):\n  arg.b'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @skip_type_union
    @refs_equal
    def test_func_annotates_union_then_uses(self):
        s = 'import a\ndef f(arg: a | 1):\n  arg.b'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_kw_default_uses_not_assigned(self):
        s = 'import a\ndef f(*_, c, b=a):\n  b'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_arg_shadows(self):
        s = 'import a\ndef f(a):\n  a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_func_decorator_uses(self):
        s = 'import a\n@a\ndef f():\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @pytest.mark.xfail(reason='Assignments are not tracked.')
    @refs_equal
    def test_func_uses_overrided_later(self):
        s = 'import a\ndef f():\n  a\na = 1\nf()'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_uses(self):
        s = 'import a\nlambda: a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_arg_shadows(self):
        s = 'import a\nlambda a: a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_arg_default_uses_not_assigned(self):
        s = 'import a\nlambda x=a: x'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_arg_default_uses_then_shadows(self):
        s = 'import a\nlambda a=a: a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_arg_shadows_used_outside(self):
        s = 'import a\nlambda a: a\na'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_lambda_kw_default_uses(self):
        s = 'import a\nlambda *b, c, d=a: 1'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @pytest.mark.xfail(reason='No reason to do this in the real world.')
    @refs_equal
    def test_global_in_outermost_scope(self):
        s = 'import a\nglobal a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_global_hits_import_after_inner_shadow(self):
        s = 'import a\ndef f():\n  a = 1\n  def g():\n    global a\n    a'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_global_skips_inner_import(self):
        s = 'a = 1\ndef f():\n  import a\n  def g():\n    global a\n    a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_global_overwritten_then_used_in_inner(self):
        s = 'import a\ndef f():\n  global a\n  a = 1\n  a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_global_overwrite_in_next_call_then_used_in_outer_before(self):
        s = 'import a\ndef f():\n  global a\n  a = 1\na\nf()'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a')]
        return s, refs

    @pytest.mark.xfail(reason='Global assigns not tracked to outer scopes.')
    @refs_equal
    def test_global_overwritten_then_used_in_outer(self):
        s = 'import a\ndef f():\n  global a\n  a = 1\nf()\na'
        refs = [('a', 'a'), ('a', 'a')]  # ref only in global statement
        return s, refs

    @refs_equal
    def test_nonlocal_hits_import_after_outer_assign(self):
        s = 'a = 1\ndef f():\n  import a\n  def g():\n    nonlocal a\n    a'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_nonlocal_skips_outer_import(self):
        s = 'import a\ndef f():\n  a = 1\n  def g():\n    nonlocal a\n    a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_nonlocal_overwritten_then_used(self):
        s = 'a = 1\ndef f():\n  import a\n  def g():\n    nonlocal a\n    a = 1\n    a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @pytest.mark.xfail(reason='Global deletes not tracked to outer scopes.')
    @refs_equal
    def test_global_deleted_then_used_in_outer(self):
        s = 'import a\ndef f():\n  global a\n  del a\nf()\na'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a')]  # refs in global and del
        return s, refs


class TestComprehension:
    @refs_equal
    def test_comp_uses_in_value(self):
        s = 'import a\n[a for b in c]'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_setcomp_uses_in_value(self):
        s = 'import a\n{a for b in c}'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_dictcomp_uses_in_value(self):
        s = 'import a\n{a: a for b in c}'
        refs = [('a', 'a'), ('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_generator_uses_in_value(self):
        s = 'import a\n(a for b in c)'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_comp_uses_in_ifs(self):
        s = 'import a\n[_ for _ in b if a]'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_comp_uses_in_iter(self):
        s = 'import a\n[_ for _ in b(a)]'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_comp_overrides(self):
        s = 'import a\n[a for a in b if a]'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_comp_overrides_used_after(self):
        s = 'import a\n[a for a in b]\na'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_multicomp_overrides(self):
        s = 'import a\n[a for a in b for b in a]'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_multicomp_uses_then_overrides(self):
        s = 'import a\n[a for b in c(a) for a in b]'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @skip_walrus
    @pytest.mark.xfail(reason='Assignments are not tracked outside of a "scope".')
    @refs_equal
    def test_comp_leaks_walrus(self):
        s = 'import a\n[a := i for i in b]\na'
        refs = [('a', 'a')]
        return s, refs


class TestClass:
    @refs_equal
    def test_class_name_shadows(self):
        s = 'import a\nclass a:\n  pass\na'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_bases_uses(self):
        s = 'import a\nclass A(a):\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_keyword_uses(self):
        s = 'import a\nclass A(kw=a):\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_starargs_uses(self):
        s = 'import a\nclass A(*a):\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_kwargs_uses(self):
        s = 'import a\nclass A(**a):\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_decorator_uses(self):
        s = 'import a\n@a\nclass A:\n  pass'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_body_pseudo_assigns(self):
        s = 'import a\nclass A:\n  a = 1\na'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_body_pseudo_shadows_for_method(self):
        s = 'import a\nclass A:\n  a = 1\n  def f(s):\n    a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_method_pseudo_shadows_inside(self):
        s = 'import a\nclass A:\n  def a(s):\n    a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_method_pseudo_shadows_after(self):
        s = 'import a\nclass A:\n  def a(s):\n    pass\na'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_method_shadows_body(self):
        s = 'import a\nclass A:\n  def a(s):\n    pass\n  a'
        refs = [('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_lambda_uses_outer(self):
        s = 'import a\nclass A:\n  b = lambda: a\na = 1'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs

    @refs_equal
    def test_class_lambda_skips_body(self):
        s = 'import a\nclass A:\n  a = 2\n  b = lambda: a'
        refs = [('a', 'a'), ('a', 'a')]
        return s, refs
