Add transform to change return None into return

Author: dflook

docs/source/transforms/remove_explicit_return_none.py

@@ -0,0 +1,5 @@
+def important(a):
+    if not a:
+        return None
+    a.adjust(1)
+    return a

docs/source/transforms/remove_explicit_return_none.rst

@@ -0,0 +1,22 @@
+Remove Explicit Return None
+===========================
+
+This transforms and ``return None`` statement into a ``return`` statement.
+A return statement with no value is equivalent to ``return None``.
+
+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.
+
+Example
+-------
+
+Input
+~~~~~
+
+.. literalinclude:: remove_explicit_return_none.py
+
+Output
+~~~~~~
+
+.. literalinclude:: remove_explicit_return_none.min.py
+    :language: python

src/python_minifier/__init__.py

@@ -23,6 +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_literal_statements import RemoveLiteralStatements
 from python_minifier.transforms.remove_object_base import RemoveObject
 from python_minifier.transforms.remove_pass import RemovePass
@@ -64,7 +65,8 @@ def minify(
     convert_posargs_to_args=True,
     preserve_shebang=True,
     remove_asserts=False,
-    remove_debug=False
+    remove_debug=False,
+    remove_explicit_return_none=True,
 ):
     """
     Minify a python module
@@ -94,6 +96,7 @@ def minify(
     :param bool preserve_shebang: Keep any shebang interpreter directive from the source in the minified output
     :param bool remove_asserts: If assert statements should be removed
     :param bool remove_debug: If conditional statements that test '__debug__ is True' should be removed
+    :param bool remove_explicit_return_none: If explicit return None statements should be replaced with a bare return
 
     :rtype: str
 
@@ -127,6 +130,9 @@ def minify(
     if remove_debug:
         module = RemoveDebug()(module)
 
+    if remove_explicit_return_none:
+        module = implicit_return_none(module)
+
     bind_names(module)
     resolve_names(module)
 

src/python_minifier/__init__.pyi

@@ -21,7 +21,8 @@ def minify(
     convert_posargs_to_args: bool = ...,
     preserve_shebang: bool = ...,
     remove_asserts: bool = ...,
-    remove_debug: bool = ...
+    remove_debug: bool = ...,
+    remove_explicit_return_none: bool = ...
 ) -> Text: ...
 
 def unparse(module: ast.Module) -> Text: ...

src/python_minifier/__main__.py

@@ -183,7 +183,12 @@ def parse_args():
         help='Remove conditional statements that test __debug__ is True',
         dest='remove_debug',
     )
-
+    minification_options.add_argument(
+        '--no-remove-explicit-return-none',
+        action='store_false',
+        help='Replace explicit return None with a bare return',
+        dest='remove_explicit_return_none',
+    )
     parser.add_argument('--version', '-v', action='version', version=version)
 
     args = parser.parse_args()
@@ -248,7 +253,8 @@ def do_minify(source, filename, minification_args):
         convert_posargs_to_args=minification_args.convert_posargs_to_args,
         preserve_shebang=minification_args.preserve_shebang,
         remove_asserts=minification_args.remove_asserts,
-        remove_debug=minification_args.remove_debug
+        remove_debug=minification_args.remove_debug,
+        remove_explicit_return_none=minification_args.remove_explicit_return_none,
     )
 
 

src/python_minifier/transforms/implicit_return_none.py

@@ -0,0 +1,15 @@
+import ast
+
+from python_minifier.util import is_ast_node
+
+
+def implicit_return_none(node):
+
+    if isinstance(node, ast.Return) and is_ast_node(node.value, 'NameConstant') and node.value.value is None:
+        # explicit return None
+        node.value = None
+
+    for child in ast.iter_child_nodes(node):
+        implicit_return_none(child)
+
+    return node

test/test_remove_explicit_return_none.py

@@ -0,0 +1,37 @@
+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
+
+
+def remove_return_none(source):
+    module = ast.parse(source, 'remove_return_none')
+
+    return implicit_return_none(module)
+
+def test_remove_return_none():
+    source = 'def a():return None'
+    expected = 'def 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():return'
+    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_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
+