Remove trailing return

Author: dflook

docs/source/transforms/remove_explicit_return_none.py

@@ -1,5 +1,7 @@
 def important(a):
-    if not a:
+    if a > 3:
+        return a
+    if a < 2:
         return None
     a.adjust(1)
-    return a
+    return None

docs/source/transforms/remove_explicit_return_none.rst

@@ -1,9 +1,11 @@
 Remove Explicit Return None
 ===========================
 
-This transforms and ``return None`` statement into a ``return`` statement.
+This transforms any ``return None`` statement into a ``return`` statement.
 A return statement with no value is equivalent to ``return None``.
 
+Removes any ``return None`` or ``return`` statements that are the last statement in a function.
+
 The transform is always safe to use and enabled by default. Disable by passing the ``remove_explicit_return_none=False`` argument to the :func:`python_minifier.minify` function,
 or passing ``--no-remove-explicit-remove-none`` to the pyminify command.
 

src/python_minifier/__init__.py

@@ -23,7 +23,7 @@
 from python_minifier.transforms.remove_annotations import RemoveAnnotations
 from python_minifier.transforms.remove_asserts import RemoveAsserts
 from python_minifier.transforms.remove_debug import RemoveDebug
-from python_minifier.transforms.implicit_return_none import implicit_return_none
+from python_minifier.transforms.remove_explicit_return_none import RemoveExplicitReturnNone
 from python_minifier.transforms.remove_literal_statements import RemoveLiteralStatements
 from python_minifier.transforms.remove_object_base import RemoveObject
 from python_minifier.transforms.remove_pass import RemovePass
@@ -131,7 +131,7 @@ def minify(
         module = RemoveDebug()(module)
 
     if remove_explicit_return_none:
-        module = implicit_return_none(module)
+        module = RemoveExplicitReturnNone()(module)
 
     bind_names(module)
     resolve_names(module)

src/python_minifier/transforms/implicit_return_none.py

@@ -0,0 +1,43 @@
+import ast
+import sys
+
+from python_minifier.transforms.suite_transformer import SuiteTransformer
+from python_minifier.util import is_ast_node
+
+class RemoveExplicitReturnNone(SuiteTransformer):
+    def __call__(self, node):
+        return self.visit(node)
+
+    def visit_Return(self, node):
+        assert isinstance(node, ast.Return)
+
+        # Remove an explicit `return None`
+
+        if sys.version_info < (3, 4) and is_ast_node(node.value, 'Name') and node.value.id == 'None':
+            # explicit return None
+            node.value = None
+
+        elif sys.version_info >= (3, 4) and is_ast_node(node.value, 'NameConstant') and node.value.value is None:
+            # explicit return None
+            node.value = None
+
+        return node
+
+    def visit_FunctionDef(self, node):
+        assert is_ast_node(node, (ast.FunctionDef, 'AsyncFunctionDef'))
+
+        node.body = [self.visit(node) for node in node.body]
+
+        # Remove an explicit `return` from the end of a function
+
+        if len(node.body) > 0 and isinstance(node.body[-1], ast.Return) and node.body[-1].value is None:
+            # explicit return None
+            node.body.pop()
+
+        if len(node.body) == 0:
+            node.body = [self.add_child(ast.Expr(value=ast.Num(0)), parent=node)]
+
+        return node
+
+    def visit_AsyncFunctionDef(self, node):
+        return self.visit_FunctionDef(node)

src/python_minifier/transforms/remove_explicit_return_none.py

@@ -1,37 +1,176 @@
 import ast
 from python_minifier import unparse
 from python_minifier.ast_compare import compare_ast
-from python_minifier.transforms.implicit_return_none import implicit_return_none
+from python_minifier.transforms.remove_explicit_return_none import RemoveExplicitReturnNone
 
 
 def remove_return_none(source):
     module = ast.parse(source, 'remove_return_none')
 
-    return implicit_return_none(module)
+    return RemoveExplicitReturnNone()(module)
 
-def test_remove_return_none():
+def test_trailing_remove_return_none():
+    source = 'def a():a=4;return None'
+    expected = 'def a():a=4'
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+
+    assert unparse(actual_ast) == expected
+
+def test_trailing_implicit_return_none():
+    source = 'def a():a=4;return'
+    expected = 'def a():a=4'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+
+def test_trailing_remove_return_none_empty_suite():
     source = 'def a():return None'
-    expected = 'def a():return'
+    expected = 'def a():0'
 
     expected_ast = ast.parse(expected)
     actual_ast = remove_return_none(source)
     compare_ast(expected_ast, actual_ast)
 
     assert unparse(actual_ast) == expected
 
-def test_implicit_return_none():
+def test_trailing_implicit_return_none_empty_suite():
     source = 'def a():return'
+    expected = 'def a():0'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+def test_trailing_return_value_unchanged():
+    source = 'def a():return 0'
     expected = source
     expected_ast = ast.parse(expected)
     actual_ast = remove_return_none(source)
     compare_ast(expected_ast, actual_ast)
     assert unparse(actual_ast) == expected
 
+def test_remove_return_none():
+    source = '''
+def a():
+    if a: return None
+    return None
+'''
+    expected = 'def a():\n\tif a:return'
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+
+    assert unparse(actual_ast) == expected
+
+def test_implicit_return_none():
+    source = '''
+def a():
+    if a: return
+    return
+'''
+    expected = 'def a():\n\tif a:return'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
 def test_return_value_unchanged():
-    source = 'def a():return 0'
+    source = '''
+def a():
+    if a: return 1
+    return 3
+'''
+    expected = 'def a():\n\tif a:return 1\n\treturn 3'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+def test_async_trailing_remove_return_none():
+    source = 'async def a():a=4;return None'
+    expected = 'async def a():a=4'
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+
+    assert unparse(actual_ast) == expected
+
+def test_async_trailing_implicit_return_none():
+    source = 'async def a():a=4;return'
+    expected = 'async def a():a=4'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+
+def test_async_trailing_remove_return_none_empty_suite():
+    source = 'async def a():return None'
+    expected = 'async def a():0'
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+
+    assert unparse(actual_ast) == expected
+
+def test_async_trailing_implicit_return_none_empty_suite():
+    source = 'async def a():return'
+    expected = 'async def a():0'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+def test_async_trailing_return_value_unchanged():
+    source = 'async def a():return 0'
     expected = source
     expected_ast = ast.parse(expected)
     actual_ast = remove_return_none(source)
     compare_ast(expected_ast, actual_ast)
     assert unparse(actual_ast) == expected
 
+def test_async_remove_return_none():
+    source = '''
+async def a():
+    if a: return None
+    return None
+'''
+    expected = 'async def a():\n\tif a:return'
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+
+    assert unparse(actual_ast) == expected
+
+def test_async_implicit_return_none():
+    source = '''
+async def a():
+    if a: return
+    return
+'''
+    expected = 'async def a():\n\tif a:return'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected
+
+def test_async_return_value_unchanged():
+    source = '''
+async def a():
+    if a: return 1
+    return 3
+'''
+    expected = 'async def a():\n\tif a:return 1\n\treturn 3'
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_return_none(source)
+    compare_ast(expected_ast, actual_ast)
+    assert unparse(actual_ast) == expected