remove_literal_statements.py

import ast

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):
        super().__init__()
        self.doc_in_module = False

    def __call__(self, node):
        self.doc_in_module = _doc_in_module(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

        node.body = self.suite(node.body, parent=node)
        return 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 suite(self, node_list, parent):
        without_literals = []

        for node in node_list:
            if self.is_literal_statement(node):
                if self.doc_in_module and is_ast_node(node.value, ast.Str):
                    node.value.value = 'doc-string stripped by python-minifier'
                    without_literals.append(node)
            else:
                without_literals.append(self.visit(node))

        if len(without_literals) == 0:
            if isinstance(parent, ast.Module):
                return []
            else:
                return [self.add_child(ast.Expr(value=ast.Num(0)), parent=parent)]

        return without_literals