Add new If guard helper methods (#1099)

* Add new if guard helper methods

Co-authored-by: Pierre Sassoulas <[email protected]>

Author: cdce8p

ChangeLog

@@ -7,11 +7,11 @@ What's New in astroid 2.7.0?
 Release date: TBA
 
 
-
 What's New in astroid 2.6.3?
 ============================
 Release date: TBA
 
+* Added ``If.is_sys_guard`` and ``If.is_typing_guard`` helper methods
 
 * Fix a bad inferenece type for yield values inside of a derived class.
 

astroid/node_classes.py

@@ -3480,6 +3480,43 @@ def _get_yield_nodes_skip_lambdas(self):
         yield from self.test._get_yield_nodes_skip_lambdas()
         yield from super()._get_yield_nodes_skip_lambdas()
 
+    def is_sys_guard(self) -> bool:
+        """Return True if IF stmt is a sys.version_info guard.
+
+        >>> node = astroid.extract_node('''
+        import sys
+        if sys.version_info > (3, 8):
+            from typing import Literal
+        else:
+            from typing_extensions import Literal
+        ''')
+        >>> node.is_sys_guard()
+        True
+        """
+        if isinstance(self.test, Compare):
+            value = self.test.left
+            if isinstance(value, Subscript):
+                value = value.value
+            if isinstance(value, Attribute) and value.as_string() == "sys.version_info":
+                return True
+
+        return False
+
+    def is_typing_guard(self) -> bool:
+        """Return True if IF stmt is a typing guard.
+
+        >>> node = astroid.extract_node('''
+        from typing import TYPE_CHECKING
+        if TYPE_CHECKING:
+            from xyz import a
+        ''')
+        >>> node.is_typing_guard()
+        True
+        """
+        return isinstance(
+            self.test, (Name, Attribute)
+        ) and self.test.as_string().endswith("TYPE_CHECKING")
+
 
 class IfExp(NodeNG):
     """Class representing an :class:`ast.IfExp` node.

tests/unittest_nodes.py

@@ -342,6 +342,64 @@ def test_block_range(self):
         self.assertEqual(self.astroid.body[1].orelse[0].block_range(7), (7, 8))
         self.assertEqual(self.astroid.body[1].orelse[0].block_range(8), (8, 8))
 
+    @staticmethod
+    def test_if_sys_guard():
+        code = builder.extract_node(
+            """
+        import sys
+        if sys.version_info > (3, 8):  #@
+            pass
+
+        if sys.version_info[:2] > (3, 8):  #@
+            pass
+
+        if sys.some_other_function > (3, 8):  #@
+            pass
+        """
+        )
+        assert isinstance(code, list) and len(code) == 3
+
+        assert isinstance(code[0], nodes.If)
+        assert code[0].is_sys_guard() is True
+        assert isinstance(code[1], nodes.If)
+        assert code[1].is_sys_guard() is True
+
+        assert isinstance(code[2], nodes.If)
+        assert code[2].is_sys_guard() is False
+
+    @staticmethod
+    def test_if_typing_guard():
+        code = builder.extract_node(
+            """
+        import typing
+        import typing as t
+        from typing import TYPE_CHECKING
+
+        if typing.TYPE_CHECKING:  #@
+            pass
+
+        if t.TYPE_CHECKING:  #@
+            pass
+
+        if TYPE_CHECKING:  #@
+            pass
+
+        if typing.SOME_OTHER_CONST:  #@
+            pass
+        """
+        )
+        assert isinstance(code, list) and len(code) == 4
+
+        assert isinstance(code[0], nodes.If)
+        assert code[0].is_typing_guard() is True
+        assert isinstance(code[1], nodes.If)
+        assert code[1].is_typing_guard() is True
+        assert isinstance(code[2], nodes.If)
+        assert code[2].is_typing_guard() is True
+
+        assert isinstance(code[3], nodes.If)
+        assert code[3].is_typing_guard() is False
+
 
 class TryExceptNodeTest(_NodeTest):
     CODE = """