Make removal literal statements more configurable

docs/source/transforms/remove_literal_statements.py

@@ -5,5 +5,13 @@
 0
 1000
 
-def test():
-    'Function docstring'
+...
+True
+False
+None
+
+class MyClass:
+    'my class docstring'
+
+    def test(self):
+        'Function docstring'

docs/source/transforms/remove_literal_statements.rst

@@ -1,14 +1,46 @@
 Remove Literal Statements
 =========================
 
-This transform removes statements that consist entirely of a literal value. This includes docstrings.
-If a statement is required, it is replaced by a literal zero expression statement.
+This transform removes statements that consist entirely of a literal value, which can include docstrings.
 
-This transform will strip docstrings from the source. If the module uses the ``__doc__`` name the module docstring will
-be retained.
+'Literal statements' are statements that consist entirely of a literal value, i.e string, bytes, number, ellipsis, True, False or None.
+These statements have no effect on the program and can be removed.
+There is one common exception to this, which is docstrings. A string literal that is the first statement in module, function or class is a docstring.
+Docstrings are made available to the program at runtime, so could affect it's behaviour if removed.
 
-This transform is disabled by default. Enable by passing the ``remove_literal_statements=True`` argument to the :func:`python_minifier.minify` function,
-or passing ``--remove-literal-statements`` to the pyminify command.
+This transform has separate options to configure removal of literal statements:
+
+    - Remove module docstring
+    - Remove function docstrings
+    - Remove class docstrings
+    - Remove any other literal statements
+
+If a literal can be removed but a statement is required, it is replaced by a literal zero expression statement.
+
+If it looks like docstrings are used by the module they will not be removed regardless of the options.
+If the module uses the ``__doc__`` name the module docstring will not be removed.
+If a ``__doc__`` attribute is used in the module, docstrings will not be removed from functions or classes.
+
+By default this transform will remove all literal statements except docstrings.
+
+Options
+-------
+
+These arguments can be used with the pyminify command:
+
+``--remove-module-docstring`` removes the module docstring if it is not used.
+
+``--remove-function-docstrings`` removes function docstrings if they are not used.
+
+``--remove-class-docstrings`` removes class docstrings if they are not used.
+
+``--no-remove-literal-expression-statements`` disables removing non-docstring literal statements.
+
+``--remove-literal-statements`` is an alias for ``--remove-module-docstring --remove-function-docstrings --remove-class-docstrings``.
+
+When using the :func:`python_minifier.minify` function you can use the  ``remove_literal_statements`` argument to control this transform.
+You can pass a boolean ``True`` to remove all literal statements (including docstrings) or a boolean ``False`` to not remove any.
+You can also pass a :class:`python_minifier.RemoveLiteralStatementsOptions` instance to specify what to remove
 
 Example
 -------

src/python_minifier/__init__.py

@@ -27,6 +27,7 @@
 from python_minifier.transforms.remove_explicit_return_none import RemoveExplicitReturnNone
 from python_minifier.transforms.remove_exception_brackets import remove_no_arg_exception_call
 from python_minifier.transforms.remove_literal_statements import RemoveLiteralStatements
+from python_minifier.transforms.remove_literal_statements_options import RemoveLiteralStatementsOptions
 from python_minifier.transforms.remove_object_base import RemoveObject
 from python_minifier.transforms.remove_pass import RemovePass
 from python_minifier.transforms.remove_posargs import remove_posargs
@@ -50,13 +51,15 @@ def __init__(self, exception, source, minified):
     def __str__(self):
         return 'Minification was unstable! Please create an issue at https://github.com/dflook/python-minifier/issues'
 
+default_remove_annotations_options = RemoveAnnotationsOptions()
+default_remove_literal_statements_options = RemoveLiteralStatementsOptions()
 
 def minify(
     source,
     filename=None,
-    remove_annotations=RemoveAnnotationsOptions(),
+    remove_annotations=default_remove_annotations_options,
     remove_pass=True,
-    remove_literal_statements=False,
+    remove_literal_statements=default_remove_literal_statements_options,
     combine_imports=True,
     hoist_literals=True,
     rename_locals=True,
@@ -87,6 +90,7 @@ def minify(
     :type remove_annotations: bool or RemoveAnnotationsOptions
     :param bool remove_pass: If Pass statements should be removed where possible
     :param bool remove_literal_statements: If statements consisting of a single literal should be removed, including docstrings
+    :type remove_literal_statements: bool or RemoveLiteralStatementsOptions
     :param bool combine_imports: Combine adjacent import statements where possible
     :param bool hoist_literals: If str and byte literals may be hoisted to the module level where possible.
     :param bool rename_locals: If local names may be shortened
@@ -114,8 +118,20 @@ def minify(
 
     add_namespace(module)
 
-    if remove_literal_statements:
-        module = RemoveLiteralStatements()(module)
+    if isinstance(remove_literal_statements, bool):
+        remove_literal_statements_options = RemoveLiteralStatementsOptions(
+            remove_module_docstring=remove_literal_statements,
+            remove_function_docstrings=remove_literal_statements,
+            remove_class_docstrings=remove_literal_statements,
+            remove_literal_expression_statements=remove_literal_statements
+        )
+    elif isinstance(remove_literal_statements, RemoveLiteralStatementsOptions):
+        remove_literal_statements_options = remove_literal_statements
+    else:
+        raise TypeError('remove_literal_statements must be a bool or RemoveLiteralStatements')
+
+    if remove_literal_statements_options:
+        module = RemoveLiteralStatements(remove_literal_statements_options)(module)
 
     if combine_imports:
         module = CombineImports()(module)

src/python_minifier/__init__.pyi

@@ -2,6 +2,7 @@ import ast
 from typing import List, Text, AnyStr, Optional, Any, Union
 
 from .transforms.remove_annotations_options import RemoveAnnotationsOptions as RemoveAnnotationsOptions
+from .transforms.remove_literal_statements_options import RemoveLiteralStatementsOptions as RemoveLiteralStatementsOptions
 
 class UnstableMinification(RuntimeError):
     def __init__(self, exception: Any, source: Any, minified: Any): ...
@@ -11,7 +12,7 @@ def minify(
     filename: Optional[str] = ...,
     remove_annotations: Union[bool, RemoveAnnotationsOptions] = ...,
     remove_pass: bool = ...,
-    remove_literal_statements: bool = ...,
+    remove_literal_statements: Union[bool, RemoveLiteralStatementsOptions] = ...,
     combine_imports: bool = ...,
     hoist_literals: bool = ...,
     rename_locals: bool = ...,
@@ -27,8 +28,10 @@ def minify(
     remove_builtin_exception_brackets: bool = ...
 ) -> Text: ...
 
+
 def unparse(module: ast.Module) -> Text: ...
 
+
 def awslambda(
     source: AnyStr,
     filename: Optional[Text] = ...,

src/python_minifier/__main__.py

@@ -8,6 +8,7 @@
 
 from python_minifier import minify
 from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
+from python_minifier.transforms.remove_literal_statements_options import RemoveLiteralStatementsOptions
 
 try:
     version = get_distribution('python_minifier').version
@@ -108,12 +109,6 @@ def parse_args():
         help='Disable removing Pass statements',
         dest='remove_pass',
     )
-    minification_options.add_argument(
-        '--remove-literal-statements',
-        action='store_true',
-        help='Enable removing statements that are just a literal (including docstrings)',
-        dest='remove_literal_statements',
-    )
     minification_options.add_argument(
         '--no-hoist-literals',
         action='store_false',
@@ -223,6 +218,38 @@ def parse_args():
         dest='remove_class_attribute_annotations',
     )
 
+    literal_statement_options = parser.add_argument_group('remove literal statements options', 'Options that affect how literal statements are removed')
+    literal_statement_options.add_argument(
+        '--remove-literal-statements',
+        action='store_true',
+        help='Enable removing all statements that are just a literal (including docstrings)',
+        dest='remove_literal_statements',
+    )
+    literal_statement_options.add_argument(
+        '--remove-module-docstring',
+        action='store_false',
+        help='Enable removing non-module docstrings',
+        dest='remove_docstrings',
+    )
+    literal_statement_options.add_argument(
+        '--remove-function-docstrings',
+        action='store_false',
+        help='Enable removing non-module docstrings',
+        dest='remove_docstrings',
+    )
+    literal_statement_options.add_argument(
+        '--remove-class-docstrings',
+        action='store_false',
+        help='Enable removing non-module docstrings',
+        dest='remove_docstrings',
+    )
+    literal_statement_options.add_argument(
+        '--no-remove-literal-expression-statements',
+        action='store_false',
+        help='Disable removing literal expression statements',
+        dest='remove_literal_expression_statements',
+    )
+
     parser.add_argument('--version', '-v', action='version', version=version)
 
     args = parser.parse_args()
@@ -275,6 +302,21 @@ def do_minify(source, filename, minification_args):
             names = [name.strip() for name in arg.split(',') if name]
             preserve_locals.extend(names)
 
+    if minification_args.remove_literal_statements is False:
+        remove_literal_statements = RemoveLiteralStatementsOptions(
+            remove_module_docstring=True,
+            remove_function_docstrings=True,
+            remove_class_docstrings=True,
+            remove_literal_expression_statements=minification_args.remove_literal_expression_statements,
+        )
+    else:
+        remove_literal_statements = RemoveLiteralStatementsOptions(
+            remove_module_docstring=minification_args.remove_module_docstring,
+            remove_function_docstrings=minification_args.remove_function_docstrings,
+            remove_class_docstrings=minification_args.remove_class_docstrings,
+            remove_literal_expression_statements=minification_args.remove_literal_expression_statements,
+        )
+
     if minification_args.remove_annotations is False:
         remove_annotations = RemoveAnnotationsOptions(
             remove_variable_annotations=False,
@@ -296,7 +338,7 @@ def do_minify(source, filename, minification_args):
         combine_imports=minification_args.combine_imports,
         remove_pass=minification_args.remove_pass,
         remove_annotations=remove_annotations,
-        remove_literal_statements=minification_args.remove_literal_statements,
+        remove_literal_statements=remove_literal_statements,
         hoist_literals=minification_args.hoist_literals,
         rename_locals=minification_args.rename_locals,
         preserve_locals=preserve_locals,

src/python_minifier/transforms/remove_literal_statements.py

@@ -1,61 +1,101 @@
 import ast
 
+from python_minifier.transforms.remove_literal_statements_options import RemoveLiteralStatementsOptions
 from python_minifier.transforms.suite_transformer import SuiteTransformer
 from python_minifier.util import is_ast_node
 
-
-def find_doc(node):
-
-    if isinstance(node, ast.Attribute):
-        if node.attr == '__doc__':
-            raise ValueError('__doc__ found!')
-
-    for child in ast.iter_child_nodes(node):
-        find_doc(child)
-
-
-def _doc_in_module(module):
-    try:
-        find_doc(module)
-        return False
-    except:
-        return True
-
-
 class RemoveLiteralStatements(SuiteTransformer):
     """
     Remove literal expressions from the code
 
     This includes docstrings
     """
 
+    def __init__(self, options):
+        assert isinstance(options, RemoveLiteralStatementsOptions)
+        self._options = options
+        super(RemoveLiteralStatements, self).__init__()
+
     def __call__(self, node):
-        if _doc_in_module(node):
-            return node
+        assert isinstance(node, ast.Module)
+
+        def has_doc_attribute(node):
+            if isinstance(node, ast.Attribute):
+                if node.attr == '__doc__':
+                    return True
+
+            for child in ast.iter_child_nodes(node):
+                if has_doc_attribute(child):
+                    return True
+
+            return False
+
+        def has_doc_binding(node):
+            for binding in node.bindings:
+                if binding.name == '__doc__':
+                    return True
+            return False
+
+        self._has_doc_attribute = has_doc_attribute(node)
+        self._has_doc_binding = has_doc_binding(node)
         return self.visit(node)
 
-    def visit_Module(self, node):
-        for binding in node.bindings:
-            if binding.name == '__doc__':
-                node.body = [self.visit(a) for a in node.body]
-                return node
+    def without_literal_statements(self, node_list):
+        """
+        Remove all literal statements except for docstrings
+        """
+
+        def is_docstring(node):
+            assert isinstance(node, ast.Expr)
+
+            if not is_ast_node(node.value, ast.Str):
+                return False
 
-        node.body = self.suite(node.body, parent=node)
-        return node
+            if is_ast_node(node.parent, (ast.FunctionDef, 'AsyncFunctionDef', ast.ClassDef, ast.Module)):
+                return node.parent.body[0] is node
 
-    def is_literal_statement(self, node):
-        if not isinstance(node, ast.Expr):
             return False
 
-        return is_ast_node(node.value, (ast.Num, ast.Str, 'NameConstant', 'Bytes'))
+        def is_literal_statement(node):
+            if not isinstance(node, ast.Expr):
+                return False
+
+            if is_docstring(node):
+                return False
+
+            return is_ast_node(node.value, (ast.Num, ast.Str, 'NameConstant', 'Bytes', ast.Ellipsis))
+
+        return [n for n in node_list if not is_literal_statement(n)]
+
+    def without_docstring(self, node_list):
+        if node_list == []:
+            return node_list
+
+        if not isinstance(node_list[0], ast.Expr):
+            return node_list
+
+        if isinstance(node_list[0].value, ast.Str):
+            return node_list[1:]
+
+        return node_list
 
     def suite(self, node_list, parent):
-        without_literals = [self.visit(n) for n in node_list if not self.is_literal_statement(n)]
+        suite = [self.visit(node) for node in node_list]
+
+        if self._options.remove_literal_expression_statements is True:
+            suite = self.without_literal_statements(node_list)
+
+        if isinstance(parent, ast.Module) and self._options.remove_module_docstring is True and self._has_doc_binding is False:
+            suite = self.without_docstring(node_list)
+        elif is_ast_node(parent, (ast.FunctionDef, 'AsyncFunctionDef')) and self._options.remove_function_docstrings is True and self._has_doc_attribute is False:
+            suite = self.without_docstring(node_list)
+        elif is_ast_node(parent, ast.ClassDef) and self._options.remove_class_docstrings is True and self._has_doc_attribute is False:
+            suite = self.without_docstring(node_list)
 
-        if len(without_literals) == 0:
+        if len(suite) == 0:
             if isinstance(parent, ast.Module):
                 return []
             else:
                 return [self.add_child(ast.Expr(value=ast.Num(0)), parent=parent)]
 
-        return without_literals
+        return suite