Skip to content

Add new If guard helper methods #1099

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 13, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
37 changes: 37 additions & 0 deletions astroid/node_classes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
58 changes: 58 additions & 0 deletions tests/unittest_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = """
Expand Down