Merge d16310d into 5315fe7

deal/linter/_extractors/__init__.py

@@ -8,6 +8,7 @@
 from .imports import get_imports
 from .prints import get_prints
 from .returns import get_returns, has_returns
+from .value import get_value
 
 
 __all__ = [
@@ -20,5 +21,6 @@
     'get_name',
     'get_prints',
     'get_returns',
+    'get_value',
     'has_returns',
 ]

deal/linter/_extractors/asserts.py

@@ -1,64 +1,19 @@
 # built-in
-import ast
 from typing import Optional
 
-# external
-import astroid
-
 # app
-from .common import TOKENS, Extractor, Token, infer
+from .common import TOKENS, Extractor, Token
+from .value import UNKNOWN, get_value
 
 
 get_asserts = Extractor()
-inner_extractor = Extractor()
 
 
 @get_asserts.register(*TOKENS.ASSERT)
 def handle_assert(expr) -> Optional[Token]:
-    # inner_extractor
-    for token in inner_extractor.handle(expr=expr.test):
-        return token
-
-    # astroid inference
-    if hasattr(expr.test, 'infer'):
-        for value in infer(expr.test):
-            if not isinstance(value, astroid.Const):
-                continue
-            if value.value:
-                continue
-            return Token(value=value.value, line=expr.lineno, col=expr.col_offset)
-    return None
-
-
-# any constant value in astroid
-@inner_extractor.register(astroid.Const)
-def handle_const(expr: astroid.Const) -> Optional[Token]:
-    if expr.value:
-        return None
-    return Token(value=expr.value, line=expr.lineno, col=expr.col_offset)
-
-
-# Python <3.8
-# string, binary string
-@inner_extractor.register(ast.Str, ast.Bytes)
-def handle_str(expr) -> Optional[Token]:  # pragma: py>=38
-    if expr.s:
-        return None
-    return Token(value=expr.s, line=expr.lineno, col=expr.col_offset)
-
-
-# Python <3.8
-# True, False, None
-@inner_extractor.register(ast.NameConstant)
-def handle_name_constant(expr: ast.NameConstant) -> Optional[Token]:  # pragma: py>=38
-    if expr.value:
+    value = get_value(expr=expr.test)
+    if value is UNKNOWN:
         return None
-    return Token(value=expr.value, line=expr.lineno, col=expr.col_offset)
-
-
-# positive number
-@inner_extractor.register(ast.Num, getattr(ast, 'Constant', None))
-def handle_num(expr) -> Optional[Token]:
-    if expr.n:
+    if value:
         return None
-    return Token(value=expr.n, line=expr.lineno, col=expr.col_offset)
+    return Token(value=value, line=expr.lineno, col=expr.col_offset + 7)

deal/linter/_extractors/returns.py

@@ -1,12 +1,9 @@
 # built-in
-import ast
 from typing import Optional
 
-# external
-import astroid
-
 # app
-from .common import TOKENS, Extractor, Token, infer, traverse
+from .common import TOKENS, Extractor, Token, traverse
+from .value import UNKNOWN, get_value
 
 
 get_returns = Extractor()
@@ -21,64 +18,16 @@ def has_returns(body: list) -> bool:
 
 
 @get_returns.register(*TOKENS.RETURN)
-def handle_returns(expr) -> Optional[Token]:
-    # inner_extractor
-    for token in inner_extractor.handle(expr=expr.value):
-        return token
-
-    # astroid inference
-    if hasattr(expr.value, 'infer'):
-        for value in infer(expr.value):
-            if isinstance(value, astroid.Const):
-                token_info = dict(line=expr.lineno, col=expr.value.col_offset)
-                return Token(value=value.value, **token_info)
-    return None
-
-
-# any constant value in astroid
-@inner_extractor.register(astroid.Const)
-def handle_const(expr: astroid.Const) -> Optional[Token]:
-    token_info = dict(line=expr.lineno, col=expr.col_offset)
-    return Token(value=expr.value, **token_info)
-
-
-# Python <3.8
-# string, binary string
-@inner_extractor.register(ast.Str, ast.Bytes)
-def handle_str(expr) -> Optional[Token]:  # pragma: py>=38
-    token_info = dict(line=expr.lineno, col=expr.col_offset)
-    return Token(value=expr.s, **token_info)
-
-
-# Python <3.8
-# True, False, None
-@inner_extractor.register(ast.NameConstant)
-def handle_name_constant(expr: ast.NameConstant) -> Optional[Token]:  # pragma: py>=38
-    token_info = dict(line=expr.lineno, col=expr.col_offset)
-    return Token(value=expr.value, **token_info)
-
-
-# positive number
-@inner_extractor.register(ast.Num, getattr(ast, 'Constant', None))
-def handle_num(expr) -> Optional[Token]:
-    token_info = dict(line=expr.lineno, col=expr.col_offset)
-    return Token(value=expr.n, **token_info)
+def handle_return(expr) -> Optional[Token]:
+    value = get_value(expr=expr.value)
+    if value is UNKNOWN:
+        return None
+    return Token(value=value, line=expr.lineno, col=expr.value.col_offset)
 
 
-# negative number
-# No need to handle astroid here, it can be inferred later.
-@inner_extractor.register(ast.UnaryOp)
-def handle_unary_op(expr: ast.UnaryOp) -> Optional[Token]:
-    # in Python 3.8 it is ast.Constant but it is subclass of ast.Num.
-    if not isinstance(expr.operand, ast.Num):
+@get_returns.register(*TOKENS.YIELD)
+def handle_yield(expr) -> Optional[Token]:
+    value = get_value(expr=expr.value)
+    if value is UNKNOWN:
         return None
-
-    token_info = dict(line=expr.lineno, col=expr.col_offset)
-    value = expr.operand.n
-    is_minus = type(expr.op) is ast.USub or expr.op == '-'
-    is_plus = type(expr.op) is ast.UAdd or expr.op == '+'
-    if is_minus:
-        value = -value
-    if is_minus or is_plus:
-        return Token(value=value, **token_info)
-    return None
+    return Token(value=value, line=expr.lineno, col=expr.value.col_offset)

deal/linter/_extractors/value.py

@@ -0,0 +1,56 @@
+# built-in
+import ast
+from contextlib import suppress
+
+# external
+import astroid
+
+# app
+from .common import Extractor, infer
+
+
+inner_extractor = Extractor()
+UNKNOWN = object()
+
+
+def get_value(expr):
+    if isinstance(expr, ast.AST):
+        with suppress(ValueError, SyntaxError):
+            return ast.literal_eval(expr)
+
+    if isinstance(expr, astroid.node_classes.NodeNG):
+        renderred = expr.as_string()
+        with suppress(ValueError, SyntaxError):
+            return ast.literal_eval(renderred)
+
+    handler = inner_extractor.handlers.get(type(expr))
+    if handler:
+        value = handler(expr=expr)
+        if value is not UNKNOWN:
+            return value
+
+    # astroid inference
+    if hasattr(expr, 'infer'):
+        for parent_expr in infer(expr=expr):
+            if parent_expr == expr:  # avoid recursion
+                continue
+            value = get_value(expr=parent_expr)
+            if value is not UNKNOWN:
+                return value
+    return UNKNOWN
+
+
+@inner_extractor.register(astroid.List, astroid.Set, astroid.Tuple)
+def handle_collections(expr):
+    result = []
+    for element_expr in expr.elts:
+        value = get_value(expr=element_expr)
+        if value is UNKNOWN:
+            return UNKNOWN
+        result.append(value)
+
+    if type(expr) is astroid.Tuple:
+        return tuple(result)
+    if type(expr) is astroid.Set:
+        return set(result)
+    return result

tests/test_linter/test_extractors/test_returns.py

@@ -20,13 +20,19 @@
     ('return b"lol"', (b'lol', )),
     ('return True', (True, )),
     ('return None', (None, )),
+    ('return [1,2,3]', ([1, 2, 3], )),
 
     ('if True: return 13', (13, )),
     ('if True:\n  return 13\nelse:\n  return 16', (13, 16)),
     ('for i in lst: return 13', (13, )),
     ('try:\n lol()\nexcept:\n 1\nelse:\n return 3', (3, )),
     ('try:\n lol()\nexcept:\n 1\nfinally:\n return 3', (3, )),
     ('with lol():\n return 3', (3, )),
+
+    ('yield 1', (1, )),
+    ('yield -1', (-1, )),
+    ('yield True', (True, )),
+    ('yield a', ()),
 ])
 def test_get_returns_simple(text, expected):
     tree = astroid.parse(text)
@@ -48,11 +54,21 @@ def test_ast_uninferrable_unary():
 
 
 @pytest.mark.parametrize('text, expected', [
-    ('return 1 + 2', (3, )),        # do a simple arithmetic
-    ('return a', ()),               # ignore uninferrable names
-    ('return ~4', (-5, )),          # handle unary bitwise NOT
-    ('a = 4\nreturn a', (4, )),    # infer variables
-    ('a = str\nreturn a', ()),     # ignore not constants
+    ('return 1 + 2', (3, )),                # do a simple arithmetic
+    ('return a', ()),                       # ignore uninferrable names
+    ('return ~4', (-5, )),                  # handle unary bitwise NOT
+
+    ('a = 4\nreturn a', (4, )),             # infer variables
+    ('a = str\nreturn a', ()),              # ignore not constants
+
+    ('a = 4\nreturn [1, a]', ([1, 4], )),   # infer variables in list
+    ('return [1, a]', ()),                  # ignore uninferralbe in list
+
+    ('a = 4\nreturn (1, a)', ((1, 4), )),   # infer variables in set
+    ('return (1, a)', ()),                  # ignore uninferralbe in set
+
+    ('a = 4\nreturn {1, a}', ({1, 4}, )),   # infer variables in tuple
+    ('return {1, a}', ()),                  # ignore uninferralbe in tuple
 ])
 def test_get_returns_inference(text, expected):
     tree = astroid.parse(text)