Fix `used-before-assignment` false positive for TYPE_CHECKING if/elif/else usage (#8071)

4 files changed, 121 insertions, 17 deletions

diff --git a/tests/functional/u/undefined/undefined_variable_py38.py b/tests/functional/u/undefined/undefined_variable_py38.py
index 2612e535f..ef774e53a 100644
--- a/tests/functional/u/undefined/undefined_variable_py38.py
+++ b/tests/functional/u/undefined/undefined_variable_py38.py
@@ -3,6 +3,7 @@
 
 # Tests for annotation of variables and potentially undefinition
 
+from typing import TYPE_CHECKING
 
 def typing_and_assignment_expression():
     """The variable gets assigned in an assignment expression"""
@@ -190,3 +191,22 @@ print(NEVER_DEFINED)  # [used-before-assignment]
 if (still_defined := False) == 1:
     NEVER_DEFINED_EITHER = 1
 print(still_defined)
+
+
+if TYPE_CHECKING:
+    import enum
+    import weakref
+elif input():
+    if input() + 1:
+        pass
+    elif (enum := None):
+        pass
+    else:
+        print(None if (weakref := '') else True)
+else:
+    pass
+
+def defined_by_walrus_in_type_checking() -> weakref:
+    """Usage of variables defined in TYPE_CHECKING blocks"""
+    print(enum)
+    return weakref
diff --git a/tests/functional/u/undefined/undefined_variable_py38.txt b/tests/functional/u/undefined/undefined_variable_py38.txt
index eb979fad2..832d8dd11 100644
--- a/tests/functional/u/undefined/undefined_variable_py38.txt
+++ b/tests/functional/u/undefined/undefined_variable_py38.txt
@@ -1,11 +1,11 @@
-used-before-assignment:17:15:17:18:typing_and_self_referencing_assignment_expression:Using variable 'var' before assignment:HIGH
-used-before-assignment:23:15:23:18:self_referencing_assignment_expression:Using variable 'var' before assignment:HIGH
-undefined-variable:48:6:48:16::Undefined variable 'no_default':UNDEFINED
-undefined-variable:56:6:56:22::Undefined variable 'again_no_default':UNDEFINED
-undefined-variable:82:6:82:19::Undefined variable 'else_assign_1':INFERENCE
-undefined-variable:105:6:105:19::Undefined variable 'else_assign_2':INFERENCE
-used-before-assignment:140:10:140:16:type_annotation_used_improperly_after_comprehension:Using variable 'my_int' before assignment:HIGH
-used-before-assignment:147:10:147:16:type_annotation_used_improperly_after_comprehension_2:Using variable 'my_int' before assignment:HIGH
-used-before-assignment:177:12:177:16:expression_in_ternary_operator_inside_container_wrong_position:Using variable 'val3' before assignment:HIGH
-used-before-assignment:181:9:181:10::Using variable 'z' before assignment:HIGH
-used-before-assignment:188:6:188:19::Using variable 'NEVER_DEFINED' before assignment:CONTROL_FLOW
+used-before-assignment:18:15:18:18:typing_and_self_referencing_assignment_expression:Using variable 'var' before assignment:HIGH
+used-before-assignment:24:15:24:18:self_referencing_assignment_expression:Using variable 'var' before assignment:HIGH
+undefined-variable:49:6:49:16::Undefined variable 'no_default':UNDEFINED
+undefined-variable:57:6:57:22::Undefined variable 'again_no_default':UNDEFINED
+undefined-variable:83:6:83:19::Undefined variable 'else_assign_1':INFERENCE
+undefined-variable:106:6:106:19::Undefined variable 'else_assign_2':INFERENCE
+used-before-assignment:141:10:141:16:type_annotation_used_improperly_after_comprehension:Using variable 'my_int' before assignment:HIGH
+used-before-assignment:148:10:148:16:type_annotation_used_improperly_after_comprehension_2:Using variable 'my_int' before assignment:HIGH
+used-before-assignment:178:12:178:16:expression_in_ternary_operator_inside_container_wrong_position:Using variable 'val3' before assignment:HIGH
+used-before-assignment:182:9:182:10::Using variable 'z' before assignment:HIGH
+used-before-assignment:189:6:189:19::Using variable 'NEVER_DEFINED' before assignment:CONTROL_FLOW
diff --git a/tests/functional/u/used/used_before_assignment_typing.py b/tests/functional/u/used/used_before_assignment_typing.py
index 21eb68799..a685bdabc 100644
--- a/tests/functional/u/used/used_before_assignment_typing.py
+++ b/tests/functional/u/used/used_before_assignment_typing.py
@@ -1,5 +1,5 @@
 """Tests for used-before-assignment for typing related issues"""
-# pylint: disable=missing-function-docstring
+# pylint: disable=missing-function-docstring,ungrouped-imports,invalid-name
 
 
 from typing import List, Optional, TYPE_CHECKING
@@ -7,8 +7,59 @@ from typing import List, Optional, TYPE_CHECKING
 if TYPE_CHECKING:
     if True:  # pylint: disable=using-constant-test
         import math
+    from urllib.request import urlopen
+    import array
+    import base64
+    import binascii
+    import bisect
+    import calendar
+    import collections
+    import copy
     import datetime
+    import email
+    import heapq
+    import json
+    import mailbox
+    import mimetypes
+    import numbers
+    import pprint
+    import types
+    import zoneinfo
+elif input():
+    import calendar, bisect  # pylint: disable=multiple-imports
+    if input() + 1:
+        import heapq
+    else:
+        import heapq
+elif input():
+    try:
+        numbers = None if input() else 1
+        import array
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        import types
+    finally:
+        copy = None
+elif input():
+    for i in range(1,2):
+        email = None
+    else:  # pylint: disable=useless-else-on-loop
+        json = None
+    while input():
+        import mailbox
+    else:  # pylint: disable=useless-else-on-loop
+        mimetypes = None
+elif input():
+    with input() as base64:
+        pass
+    with input() as temp:
+        import binascii
+else:
     from urllib.request import urlopen
+    zoneinfo: str = ''
+    def pprint():
+        pass
+    class collections:  # pylint: disable=too-few-public-methods,missing-class-docstring
+        pass
 
 class MyClass:
     """Type annotation or default values for first level methods can't refer to their own class"""
@@ -111,3 +162,36 @@ class ConditionalImportGuardedWhenUsed:  # pylint: disable=too-few-public-method
     """Conditional imports also guarded by TYPE_CHECKING when used."""
     if TYPE_CHECKING:
         print(urlopen)
+
+
+class TypeCheckingMultiBranch:  # pylint: disable=too-few-public-methods,unused-variable
+    """Test for defines in TYPE_CHECKING if/elif/else branching"""
+    def defined_in_elif_branch(self) -> calendar.Calendar:
+        print(bisect)
+        return calendar.Calendar()
+
+    def defined_in_else_branch(self) -> urlopen:
+        print(zoneinfo)
+        print(pprint())
+        print(collections())
+        return urlopen
+
+    def defined_in_nested_if_else(self) -> heapq:
+        print(heapq)
+        return heapq
+
+    def defined_in_try_except(self) -> array:
+        print(types)
+        print(copy)
+        print(numbers)
+        return array
+
+    def defined_in_loops(self) -> json:
+        print(email)
+        print(mailbox)
+        print(mimetypes)
+        return json
+
+    def defined_in_with(self) -> base64:
+        print(binascii)
+        return base64
diff --git a/tests/functional/u/used/used_before_assignment_typing.txt b/tests/functional/u/used/used_before_assignment_typing.txt
index a327a7a73..c0a31fae0 100644
--- a/tests/functional/u/used/used_before_assignment_typing.txt
+++ b/tests/functional/u/used/used_before_assignment_typing.txt
@@ -1,5 +1,5 @@
-undefined-variable:17:21:17:28:MyClass.incorrect_typing_method:Undefined variable 'MyClass':UNDEFINED
-undefined-variable:22:26:22:33:MyClass.incorrect_nested_typing_method:Undefined variable 'MyClass':UNDEFINED
-undefined-variable:27:20:27:27:MyClass.incorrect_default_method:Undefined variable 'MyClass':UNDEFINED
-used-before-assignment:88:35:88:39:MyFourthClass.is_close:Using variable 'math' before assignment:HIGH
-used-before-assignment:101:20:101:28:VariableAnnotationsGuardedByTypeChecking:Using variable 'datetime' before assignment:HIGH
+undefined-variable:68:21:68:28:MyClass.incorrect_typing_method:Undefined variable 'MyClass':UNDEFINED
+undefined-variable:73:26:73:33:MyClass.incorrect_nested_typing_method:Undefined variable 'MyClass':UNDEFINED
+undefined-variable:78:20:78:27:MyClass.incorrect_default_method:Undefined variable 'MyClass':UNDEFINED
+used-before-assignment:139:35:139:39:MyFourthClass.is_close:Using variable 'math' before assignment:HIGH
+used-before-assignment:152:20:152:28:VariableAnnotationsGuardedByTypeChecking:Using variable 'datetime' before assignment:HIGH