drift/mycodemod/replace_functions.py

import argparse
from ast import Expression, literal_eval
from typing import Union

import libcst as cst
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
from libcst.codemod.visitors import AddImportsVisitor


class ReplaceFunctionCommand(VisitorBasedCodemodCommand):

    # Add a description so that future codemodders can see what this does.
    DESCRIPTION: str = "Replaces the body of a function with pass."


    def __init__(self, context: CodemodContext) -> None:
        # Initialize the base class with context, and save our args. Remember, the
        # "dest" for each argument we added above must match a parameter name in
        # this init.
        super().__init__(context)


    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
        functions_docstring = updated_node.get_docstring()
        docstring_should_be = '"""No docstring here yet."""'
        if functions_docstring is not None:
            docstring_should_be = '"""\n{}\n\n"""'.format(functions_docstring)

        replace_function = cst.FunctionDef(
            name=updated_node.name, 
            params=updated_node.params,#cst.Parameters(), 
            body=cst.IndentedBlock(
                body=[
                    cst.SimpleStatementLine(
                        body=[
                            cst.Expr(
                                value=cst.SimpleString(
                                    value=docstring_should_be,
                                    lpar=[],
                                    rpar=[],
                                ),
                                semicolon=cst.MaybeSentinel.DEFAULT,
                            ),
                        ],
                        leading_lines=[],
                        trailing_whitespace=cst.TrailingWhitespace(
                        whitespace=cst.SimpleWhitespace(
                            value='',
                            ),
                            comment=None,
                            newline=cst.Newline(
                                value=None,
                        ),
            ),
                    ),
                    cst.SimpleStatementLine(
                                body=[
                                    cst.Pass(),
                                ],
                            ),
                    ] 
            )
        )
        
        return replace_function

        

    
    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
        new_body = []
        for body_item in updated_node.body.body:
            if type(body_item) is cst.FunctionDef:
                new_body.append(self.leave_FunctionDef(body_item, body_item))
        
        return updated_node.with_changes(
            body=cst.IndentedBlock(new_body)
        )
    
    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
        new_module_body=[]
        for node in original_node.body:
            if type(node) is cst.FunctionDef:
                new_module_body.append(self.leave_FunctionDef(node,node))
            
            if type(node) is cst.ClassDef:
                new_module_body.append(self.leave_ClassDef(node,node))

        replace_function = cst.Module(
            body= new_module_body
        )
        
        return replace_function