AST renaming modules

file_ast.py

import ast

# Replace the old package name with the new package name
def replace_imports(node, old_name, new_name):
    if isinstance(node, ast.ImportFrom):
        if node.module == old_name:
            node.module = new_name
        for alias in node.names:
            if alias.name == old_name:
                alias.name = new_name
    elif isinstance(node, ast.Import):
        for alias in node.names:
            if alias.name == old_name:
                alias.name = new_name

# Create a new import package based on the module name
def create_import(node, new_name):
    if isinstance(node, ast.Module):
        body = node.body
        for i in range(len(body)):
            if isinstance(body[i], ast.ImportFrom):
                if body[i].module.startswith(new_name):
                    continue
                body.insert(i, ast.parse(f"from {new_name} import {body[i].module.split('.')[-1]}").body[0])

# Parse the input Python script and modify it
def modify_imports(source, old_name, new_name):
    tree = ast.parse(source)
    for node in ast.walk(tree):
        replace_imports(node, old_name, new_name)
        create_import(node, new_name)
    return ast.unparse(tree)

# Example usage
if __name__ == '__main__':
    old_name = 'numpy'
    new_name = 'np'
    with open('example_script.py', 'r') as f:
        source = f.read()
    modified_source = modify_imports(source, old_name, new_name)
    print(modified_source)

file_tokenize.py

import tokenize
import io

def replace_imports(file_path, old_module, new_module):
    with open(file_path, 'rb') as f:
        tokens = tokenize.generate_tokens(f.readline)
        new_content = io.StringIO()
        for toknum, tokval, _, _, _ in tokens:
            if toknum == tokenize.NAME and tokval == 'import':
                # If the next token is the old module name, replace it with the new module name
                next_tok = next(tokens)
                if next_tok[0] == tokenize.NAME and next_tok[1] == old_module:
                    new_content.write('import ' + new_module)
                else:
                    new_content.write(tokval + ' ' + next_tok[1])
            elif toknum == tokenize.NAME and tokval == 'from':
                # If the next token is the old module name, replace it with the new module name
                next_tok = next(tokens)
                if next_tok[0] == tokenize.NAME and next_tok[1] == old_module:
                    new_content.write('from ' + new_module)
                    # Copy the rest of the "from" statement
                    while next_tok[1] != 'import':
                        next_tok = next(tokens)
                        new_content.write(' ' + next_tok[1])
                    new_content.write(' import')
                else:
                    new_content.write(tokval + ' ' + next_tok[1])
            else:
                new_content.write(tokval)
        new_content.seek(0)
        with open(file_path, 'w') as f:
            f.write(new_content.read())

# Example usage:
replace_imports('example.py', 'numpy', 'np')