From 93fbd668368cac7bdd58d4c86287284541847574 Mon Sep 17 00:00:00 2001
From: Libba Lawrence <llawrence@microsoft.com>
Date: Wed, 17 Apr 2024 12:30:20 -0700
Subject: [PATCH] [Pylint] bug fix typing (#8104)

* recognize else

* check string of the if

* update

* pass

* Update tools/pylint-extensions/azure-pylint-guidelines-checker/CHANGELOG.md
---
 .../CHANGELOG.md                              |  4 +++
 .../pylint_guidelines_checker.py              | 31 ++++++++++++-------
 .../azure-pylint-guidelines-checker/setup.py  |  2 +-
 .../tests/test_pylint_custom_plugins.py       | 19 ++++++++++++
 4 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/CHANGELOG.md b/tools/pylint-extensions/azure-pylint-guidelines-checker/CHANGELOG.md
index 523606fb676..65beb3a3bca 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/CHANGELOG.md
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Release History
 
+## 0.4.1 (2024-04-17)
+
+- Bug fix for typing under TYPE_CHECKING block.
+
 ## 0.4.0 (2024-04-15)
 
 - Checker to enforce no importing typing under TYPE_CHECKING block.
diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py b/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
index 3bba4819fb0..f79af8f887a 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
@@ -2725,23 +2725,30 @@ class NoImportTypingFromTypeCheck(BaseChecker):
 
     def visit_importfrom(self, node):
         """Check that we aren't importing from typing under if TYPE_CHECKING."""
-        if isinstance(node.parent, astroid.If) and (node.modname == "typing" or node.modname == "typing_extensions"):
-            self.add_message(
-                msgid=f"no-typing-import-in-type-check",
-                node=node,
-                confidence=None,
-            )
-
-    def visit_import(self, node):
-        """Check that we aren't importing from typing under if TYPE_CHECKING."""
-        if isinstance(node.parent, astroid.If):
-            for name, _ in node.names:
-                if name == "typing" or name == "typing_extensions":
+        try:
+            if isinstance(node.parent, astroid.If) and "TYPE_CHECKING" in node.parent.as_string():
+                if node.modname == "typing" or node.modname == "typing_extensions":
                     self.add_message(
                         msgid=f"no-typing-import-in-type-check",
                         node=node,
                         confidence=None,
                     )
+        except:
+            pass
+
+    def visit_import(self, node):
+        """Check that we aren't importing from typing under if TYPE_CHECKING."""
+        try:
+            if isinstance(node.parent, astroid.If) and "TYPE_CHECKING" in node.parent.as_string():
+                for name, _ in node.names:
+                    if name == "typing" or name == "typing_extensions":
+                        self.add_message(
+                            msgid=f"no-typing-import-in-type-check",
+                            node=node,
+                            confidence=None,
+                        )
+        except:
+            pass
 
 # if a linter is registered in this function then it will be checked with pylint
 def register(linter):
diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/setup.py b/tools/pylint-extensions/azure-pylint-guidelines-checker/setup.py
index 5e8cb4aaea7..12db0fc7bdb 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/setup.py
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/setup.py
@@ -6,7 +6,7 @@
 
 setup(
     name="azure-pylint-guidelines-checker",
-    version="0.4.0",
+    version="0.4.1",
     url="http://github.com/Azure/azure-sdk-for-python",
     license="MIT License",
     description="A pylint plugin which enforces azure sdk guidelines.",
diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py b/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
index d2a2f21261d..dfdca1775bc 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
@@ -4811,3 +4811,22 @@ def test_allowed_imports(self):
         )
         with self.assertNoMessages():
             self.checker.visit_importfrom(importfrom_node)
+
+    def test_allowed_import_else(self):
+        """Check that illegal imports raise warnings"""
+        ima, imb, imc, imd = astroid.extract_node(
+            """
+            if sys.version_info >= (3, 9):
+                from collections.abc import MutableMapping
+            else:
+                from typing import MutableMapping #@
+                import typing #@
+                import typing_extensions #@
+                from typing_extensions import Protocol #@
+            """
+        )
+        with self.assertNoMessages():
+            self.checker.visit_importfrom(ima)
+            self.checker.visit_import(imb)
+            self.checker.visit_import(imc)
+            self.checker.visit_importfrom(imd)
\ No newline at end of file